dask.array.argtopk

dask.array.argtopk

dask.array.argtopk(a, k, axis=- 1, split_every=None)[source]

Extract the indices of the k largest elements from a on the given axis, and return them sorted from largest to smallest. If k is negative, extract the indices of the -k smallest elements instead, and return them sorted from smallest to largest.

This performs best when k is much smaller than the chunk size. All results will be returned in a single chunk along the given axis.

Parameters
x: Array

Data being sorted

k: int
axis: int, optional
split_every: int >=2, optional

See topk(). The performance considerations for topk also apply here.

Returns
Selection of np.intp indices of x with size abs(k) along the given axis.

Examples

>>> import dask.array as da
>>> x = np.array([5, 1, 3, 6])
>>> d = da.from_array(x, chunks=2)
>>> d.argtopk(2).compute()
array([3, 0])
>>> d.argtopk(-2).compute()
array([1, 2])