dpnp.argpartition

dpnp.argpartition(a, kth, axis=-1, kind='introselect', order=None)

Perform an indirect partition along the given axis using the algorithm specified by the kind keyword. It returns an array of indices of the same shape as a that index data along the given axis in partitioned order.

For full documentation refer to numpy.argpartition.

Parameters:

a{dpnp.ndarray, usm_ndarray}

Array to sort.

kth{int, sequence of ints}

Element index to partition by. The k-th element will be in its final sorted position and all smaller elements will be moved before it and all larger elements behind it. The order of all elements in the partitions is undefined. If provided with a sequence of k-th it will partition all of them into their sorted position at once.

axis{None, int}, optional

Axis along which to sort. If None, the array is flattened before sorting. The default is -1, which sorts along the last axis.

Default: -1.

Returns:

index_arraydpnp.ndarray

Array of indices that partition a along the specified axis. If a is one-dimensional, a[index_array] yields a partitioned a. More generally, dpnp.take_along_axis(a, index_array, axis=axis) always yields the partitioned a, irrespective of dimensionality.

Limitations

Parameters kind and order are only supported with its default value. Otherwise NotImplementedError exception will be raised.

See also

dpnp.partition

Describes partition algorithms used.

dpnp.ndarray.partition

Inplace partition.

dpnp.argsort

Full indirect sort.

dpnp.take_along_axis

Apply index_array from argpartition to an array as if by calling partition.

Examples

One dimensional array:

>>> import dpnp as np
>>> x = np.array([3, 4, 2, 1])
>>> x[np.argpartition(x, 3)]
array([2, 1, 3, 4]) # may vary
>>> x[np.argpartition(x, (1, 3))]
array([1, 2, 3, 4]) # may vary

Multi-dimensional array:

>>> x = np.array([[3, 4, 2], [1, 3, 1]])
>>> index_array = np.argpartition(x, kth=1, axis=-1)
>>> # below is the same as np.partition(x, kth=1)
>>> np.take_along_axis(x, index_array, axis=-1)
array([[2, 3, 4],
       [1, 1, 3]])