array_api_extra.argpartition

array_api_extra.argpartition(a, kth, /, axis=-1, *, xp=None)

Perform an indirect partition along the given axis.

It returns an array of indices of the same shape as a that index data along the given axis in partitioned order.

Parameters:

  • a (Array) – Input array.

  • kth (int) – Element index to partition by.

  • axis (int, optional) – Axis along which to partition. The default is -1 (the last axis). If None, the flattened array is used.

  • xp (array_namespace, optional) – The standard-compatible namespace for x. Default: infer.

Returns:

Array of indices that partition a along the specified axis.

Return type:

index_array

Notes

If xp implements argpartition or an equivalent function e.g. topk for torch), complexity will likely be O(n). If not, this function simply calls xp.argsort and complexity is O(n log n).