EHN: add numpy.topk

quarrying · quarrying · commit 3ecb4e4c32c0 · 2021-05-28T20:34:47.000+08:00
diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
@@ -23,7 +23,7 @@
     'ndim', 'nonzero', 'partition', 'prod', 'product', 'ptp', 'put',
     'ravel', 'repeat', 'reshape', 'resize', 'round_',
     'searchsorted', 'shape', 'size', 'sometrue', 'sort', 'squeeze',
-    'std', 'sum', 'swapaxes', 'take', 'trace', 'transpose', 'var',
+    'std', 'sum', 'swapaxes', 'take', 'topk', 'trace', 'transpose', 'var',
 ]
 
 _gentype = types.GeneratorType
@@ -1276,6 +1276,109 @@ def argmin(a, axis=None, out=None):
     return _wrapfunc(a, 'argmin', axis=axis, out=out)
 
 
+def _topk_dispatcher(a, k, axis=None, largest=None, sorted=None):
+    return (a,)
+
+
+@array_function_dispatch(_topk_dispatcher)
+def topk(a, k, axis=-1, largest=True, sorted=True):
+    """
+    Finds values and indices of the `k` largest/smallest 
+    elements along the given `axis`.
+
+    Parameters
+    ----------
+    a: array_like
+        Array with given axis at least k.
+    k: int
+        Number of top elements to look for along the given axis.
+    axis: int or None, optional
+        Axis along which to find topk. If None, the array is flattened 
+        before sorting. The default is -1 (the last axis).
+    largest: bool, optional
+        Controls whether to return largest or smallest elements. 
+    sorted: bool, optional
+        If true the resulting k elements will be sorted by the values.
+
+    Returns
+    -------
+    topk_values : ndarray
+        Array of values of `k` largest/smallest elements 
+        along the specified `axis`.
+    topk_indices: ndarray, int
+        Array of indices of `k` largest/smallest elements 
+        along the specified `axis`.
+
+    See Also
+    --------
+    sort : Describes sorting algorithms used.
+    argsort : Indirect sort.
+    partition : Describes partition algorithms used.
+    argpartition : Indirect partial sort.
+    take_along_axis : Take values from the input array by 
+                      matching 1d index and data slices.
+
+    Examples
+    --------
+    One dimensional array:
+
+    >>> x = np.array([3, 1, 2])
+    >>> np.topk(x, 2)
+    (array([3, 2]), array([0, 2], dtype=int64))
+
+    Two-dimensional array:
+
+    >>> x = np.array([[0, 3, 4], [2, 2, 1], [5, 1, 2]])
+    >>> val, ind = np.topk(x, 2, axis=1) # along the last axis
+    >>> val
+    array([[4, 3],
+           [2, 2],
+           [5, 2]])
+    >>> ind
+    array([[2, 1],
+           [0, 1],
+           [0, 2]])
+
+    >>> val, ind = np.topk(x, 2, axis=None) # along the flattened array
+    >>> val
+    array([5, 4])
+    >>> ind
+    array([6, 2])
+
+    >>> val, ind = np.topk(x, 2, axis=0) # along the first axis
+    >>> val
+    array([[5, 3, 4],
+           [2, 2, 2]])
+    >>> ind
+    array([[2, 0, 0],
+           [1, 1, 2]])
+    """
+    if axis is None:
+        axis_size = a.size
+    else:
+        axis_size = a.shape[axis]
+    assert 1 <= k <= axis_size
+
+    a = np.asanyarray(a)
+    if largest:
+        index_array = np.argpartition(a, axis_size-k, axis=axis)
+        topk_indices = np.take(index_array, -np.arange(k)-1, axis=axis)
+    else:
+        index_array = np.argpartition(a, k-1, axis=axis)
+        topk_indices = np.take(index_array, np.arange(k), axis=axis)
+    topk_values = np.take_along_axis(a, topk_indices, axis=axis)
+    if sorted:
+        sorted_indices_in_topk = np.argsort(topk_values, axis=axis)
+        if largest:
+            sorted_indices_in_topk = np.flip(sorted_indices_in_topk, axis=axis)
+        sorted_topk_values = np.take_along_axis(
+            topk_values, sorted_indices_in_topk, axis=axis)
+        sorted_topk_indices = np.take_along_axis(
+            topk_indices, sorted_indices_in_topk, axis=axis)
+        return sorted_topk_values, sorted_topk_indices
+    return topk_values, topk_indices
+
+
 def _searchsorted_dispatcher(a, v, side=None, sorter=None):
     return (a, v, sorter)
 
diff --git a/numpy/core/fromnumeric.pyi b/numpy/core/fromnumeric.pyi
@@ -151,6 +151,23 @@ def argmin(
     out: Optional[ndarray] = ...,
 ) -> Any: ...
 
+@overload
+def topk(
+    a: ArrayLike,
+    k: int = ...,
+    axis: None = ...,
+    largest: Optional[bool] = ...,
+    sorted: Optional[bool] = ...,
+) -> Tuple[ndarray, ndarray]: ...
+@overload
+def topk(
+    a: ArrayLike,
+    k: int = ...,
+    axis: Optional[int] = ...,
+    largest: Optional[bool] = ...,
+    sorted: Optional[bool] = ...,
+) -> Tuple[ndarray, ndarray]: ...
+
 @overload
 def searchsorted(
     a: ArrayLike,
