numpy.top_k#

numpy.top_k(a, k, /, *, axis=-1, mode='largest', sorted=True)[source]#

Returns the k largest or smallest elements and their indices along an axis.

A tuple of (values, indices) is returned, where values and indices are the values and indices, respectively, of the largest/smallest elements of each row of the input array in the given axis.

Parameters:

a: array_like

The source array

k: int

The number of largest/smallest elements to return. k must be a non-negative integer and within indexable range specified by axis.

axis: int, optional

Axis along which to find the largest/smallest elements. The default is -1 (the last axis).

mode: {“largest”, “smallest”}, optional

If “largest”, the largest elements are returned. If “smallest”, the smallest elements are returned. The default is “largest”.

Similarly to sorts, NaN values are pushed to the end and therefore only present in the output if they are among the top k values, regardless of the value of mode.

sorted: bool, optional

If True, the top k elements are returned in sorted order. If False, sorted order is not guaranteed. The default is True.

Returns:

tuple_of_array: tuple: The output tuple of (topk_values, topk_indices), where topk_values are the top k values and topk_indices are the corresponding indices. Both arrays are of the shape of the input array with the dimension along axis replaced by k.

See also

argpartition: Indirect partition.
sort: Full sorting.

Notes

The returned indices are not guaranteed to be stable, i.e., the order of the returned indices for any duplicate values is not guaranteed to be the same as their order in the input array. This is the case regardless of the value of the sorted parameter.

Examples

>>> a = np.array([[1,2,3,4,5], [5,4,3,2,1]])
>>> np.top_k(a, 2)
(array([[5, 4],
        [5, 4]]),
 array([[4, 3],
        [0, 1]]))
>>> np.top_k(a, 2, axis=0)
(array([[5, 4, 3, 4, 5],
       [1, 2, 3, 2, 1]]),
 array([[1, 1, 0, 0, 0],
       [0, 0, 1, 1, 1]]))
>>> np.top_k(a, 2, axis=1, mode="smallest")
(array([[1, 2],
        [1, 2]]),
 array([[0, 1],
        [4, 3]]))
>>> np.top_k(np.array([1., 2., 3., np.nan]), 2)
(array([3., 2.]), array([2, 1]))