searchsorted¶

array_api_extra.searchsorted(x1, x2, /, *, side='left', xp=None)¶

Find indices where elements should be inserted to maintain order.

Find the indices into a sorted array x1 such that if the elements in x2 were inserted before the indices, the resulting array would remain sorted.

The behavior of this function is similar to that of array_api.searchsorted(), but it relaxes the requirement that x1 must be one-dimensional. This function is vectorized, treating slices along the last axis as elements and preceding axes as batch (or “loop”) dimensions.

Parameters:

x1 (object) – Input array. Should have a real-valued data type. Must be sorted in ascending order along the last axis.
x2 (object) – Array containing search values. Should have a real-valued data type. Must have the same shape as x1 except along the last axis.
side (Literal['left', 'right']) – Argument controlling which index is returned if an element of x2 is equal to one or more elements of x1: 'left' returns the index of the first of these elements; 'right' returns the next index after the last of these elements. Default: 'left'.
xp (ModuleType | None) – The standard-compatible namespace for the array arguments. Default: infer.

Returns:

Array – An array of indices with the same shape as x2.

Return type:

object

Examples

>>> import array_api_strict as xp
>>> import array_api_extra as xpx
>>> x = xp.asarray([11, 12, 13, 13, 14, 15])
>>> xpx.searchsorted(x, xp.asarray([10, 11.5, 14.5, 16]), xp=xp)
Array([0, 1, 5, 6], dtype=array_api_strict.int64)
>>> xpx.searchsorted(x, xp.asarray(13), xp=xp)
Array(2, dtype=array_api_strict.int64)
>>> xpx.searchsorted(x, xp.asarray(13), side='right', xp=xp)
Array(4, dtype=array_api_strict.int64)

searchsorted is vectorized along the last axis.

>>> x1 = xp.asarray([[1., 2., 3., 4.], [5., 6., 7., 8.]])
>>> x2 = xp.asarray([[1.1, 3.3], [6.6, 8.8]])
>>> xpx.searchsorted(x1, x2, xp=xp)
Array([[1, 3],
       [2, 4]], dtype=array_api_strict.int64)