Python numpy argpartition用法及代码示例

本文简要介绍 python 语言中 numpy.argpartition 的用法。

用法:

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

使用 kind 关键字指定的算法沿给定轴执行间接分区。它以分区顺序沿给定轴返回与该索引数据具有相同形状的索引数组。

参数：

a： array_like

要排序的数组。

kth： int 或整数序列

要分区的元素索引。 k-th 元素将处于其最终排序位置，所有较小的元素都将移到它之前，所有较大的元素都将移到它后面。分区中所有元素的顺序未定义。如果提供了 k-th 的序列，它将立即将所有这些分区到它们的排序位置。

axis： int 或无，可选

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

kind： {‘introselect’}，可选

选择算法。默认为‘introselect’

order： str 或 str 列表，可选

当 a 是定义了字段的数组时，该参数指定首先比较哪些字段，第二个等。单个字段可以指定为字符串，不需要指定所有字段，但仍会使用未指定的字段，在他们在 dtype 中出现的顺序，以打破关系。

返回：

index_array： ndarray，int

分区索引数组a沿着指定的轴。如果a是一维的，a[index_array]产生一个分区a.更普遍，np.take_along_axis(a, index_array, axis=a)总是产生分区的a，与维度无关。

注意：

有关不同选择算法的说明，请参见 partition 。

例子：

一维数组：

>>> x = np.array([3, 4, 2, 1])
>>> x[np.argpartition(x, 3)]
array([2, 1, 3, 4])
>>> x[np.argpartition(x, (1, 3))]
array([1, 2, 3, 4])

>>> x = [3, 4, 2, 1]
>>> np.array(x)[np.argpartition(x, 3)]
array([2, 1, 3, 4])

多维数组：

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