diff --git a/spec/draft/API_specification/searching_functions.rst b/spec/draft/API_specification/searching_functions.rst
index 01ab4e82a..c952f1aad 100644
--- a/spec/draft/API_specification/searching_functions.rst
+++ b/spec/draft/API_specification/searching_functions.rst
@@ -23,4 +23,5 @@ Objects in API
argmax
argmin
nonzero
+ searchsorted
where
diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py
index e586a7656..dda000e74 100644
--- a/src/array_api_stubs/_draft/searching_functions.py
+++ b/src/array_api_stubs/_draft/searching_functions.py
@@ -1,7 +1,7 @@
-__all__ = ["argmax", "argmin", "nonzero", "where"]
+__all__ = ["argmax", "argmin", "nonzero", "searchsorted", "where"]
-from ._types import Optional, Tuple, array
+from ._types import Optional, Tuple, Literal, array
def argmax(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> array:
@@ -87,6 +87,56 @@ def nonzero(x: array, /) -> Tuple[array, ...]:
"""
+def searchsorted(
+ x1: array,
+ x2: array,
+ /,
+ *,
+ side: Literal["left", "right"] = "left",
+ sorter: Optional[array] = None,
+) -> array:
+ """
+ Finds the indices into ``x1`` such that, if the corresponding elements in ``x2`` were inserted before the indices, the order of ``x1``, when sorted in ascending order, would be preserved.
+
+ Parameters
+ ----------
+ x1: array
+ input array. Must be a one-dimensional array. Should have a real-valued data type. If ``sorter`` is ``None``, must be sorted in ascending order; otherwise, ``sorter`` must be an array of indices that sort ``x1`` in ascending order.
+ x2: array
+ array containing search values. Should have a real-valued data type.
+ side: Literal['left', 'right']
+ argument controlling which index is returned if a value lands exactly on an edge.
+
+ Let ``x`` be an array of rank ``N`` where ``v`` is an individual element given by ``v = x2[n,m,...,j]``.
+
+ If ``side == 'left'``, then
+
+ - each returned index ``i`` must satisfy the index condition ``x1[i-1] < v <= x1[i]``.
+ - if no index satisfies the index condition, then the returned index for that element must be ``0``.
+
+ Otherwise, if ``side == 'right'``, then
+
+ - each returned index ``i`` must satisfy the index condition ``x1[i-1] <= v < x1[i]``.
+ - if no index satisfies the index condition, then the returned index for that element must be ``N``, where ``N`` is the number of elements in ``x1``.
+
+ Default: ``'left'``.
+ sorter: Optional[array]
+ array of indices that sort ``x1`` in ascending order. The array must have the same shape as ``x1`` and have an integer data type. Default: ``None``.
+
+ Returns
+ -------
+ out: array
+ an array of indices with the same shape as ``x2``. The returned array must have the default array index data type.
+
+ Notes
+ -----
+
+ For real-valued floating-point arrays, the sort order of NaNs and signed zeros is unspecified and thus implementation-dependent. Accordingly, when a real-valued floating-point array contains NaNs and signed zeros, what constitutes ascending order may vary among specification-conforming array libraries.
+
+ While behavior for arrays containing NaNs and signed zeros is implementation-dependent, specification-conforming libraries should, however, ensure consistency with ``sort`` and ``argsort`` (i.e., if a value in ``x2`` is inserted into ``x1`` according to the corresponding index in the output array and ``sort`` is invoked on the resultant array, the sorted result should be an array in the same order).
+ """
+
+
def where(condition: array, x1: array, x2: array, /) -> array:
"""
Returns elements chosen from ``x1`` or ``x2`` depending on ``condition``.
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy