numpy.argpartition#

numpy.argpartition(a, kth, axis=-1, kind='introselect', order=None)[source]#

使用 kind 关键字指定的算法,沿给定轴执行间接分区。它返回一个与 a 形状相同的索引数组,这些索引按分区顺序对给定轴上的数据进行索引。

参数:
a类数组对象

要排序的数组。

kth整数或整数序列

用于分区的元素索引。第 k 个元素将处于其最终排序位置,所有较小的元素都将被移动到它之前,所有较大的元素都将被移动到它之后。分区中所有元素的顺序是不确定的。如果提供了一个 k-th 序列,它将一次性将它们全部划分到各自的排序位置。

axisint 或 None, 可选

进行排序的轴。默认值为 -1(最后一个轴)。如果为 None,则使用展平后的数组。

kind{‘introselect’}, 可选

选择算法。默认值为 ‘introselect’

orderstr 或 str 列表, optional

a 是一个具有已定义字段的数组时,此参数指定首先、其次等要比较的字段。单个字段可以指定为字符串,并且不必指定所有字段,但未指定的字段仍将按照它们在 dtype 中出现的顺序用于打破平局。

返回:
index_arrayndarray, int

沿指定轴对 a 进行分区的索引数组。如果 a 是一维的,a[index_array] 将产生一个分区后的 a。更通用地说,np.take_along_axis(a, index_array, axis=axis) 总是产生分区后的 a,无论维度如何。

另请参阅

partition

描述所使用的分区算法。

ndarray.partition

原地分区。

argsort

完全间接排序。

take_along_axis

应用 argpartition 中的 index_array 到数组,效果如同调用 partition。

备注

返回的索引不保证根据值进行排序。此外,默认的选择算法 introselect 是不稳定的,因此返回的索引不保证是该元素最早/最晚出现的位置。

argpartition 适用于包含 nan 值的实数/复数输入,关于增强排序顺序和不同选择算法的说明,请参见 partition

示例

一维数组

>>> import numpy 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
>>> x = [3, 4, 2, 1]
>>> np.array(x)[np.argpartition(x, 3)]
array([2, 1, 3, 4]) # may vary

多维数组

>>> 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]])