NumPy的实用函数整理之argpartition

NumPy

NumPy是使用Python进行科学计算的基本包。它包括以下内容:

• 强大的n维数组对象
• 复杂的(广播)函数
• 集成C/ C++和Fortran代码的工具
• 有用的线性代数、傅里叶变换和随机数功能

除了其明显的科学用途外，NumPy还可以用作通用数据的高效多维容器。可以定义任意数据类型。这使得NumPy可以无缝且快速地与各种数据库集成。

argpartition()

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

其中参数：

• a是numpy数组
• kth是从小到大第kth

argpartition()函数划分重组数组,返回的是重组后数据的索引数组。这个可以很快地找出第 kth大的数的位置，以及大于 kth（排在kth后面）数的位置,一般numpy中以arg开头的函数都是返回下标，而不改变原数组。下面具体看两种类型的例子：

输入：

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

输入：

x[np.argpartition(x, -4)[-4:]] 

输出：

array([1, 4, 2, 3])

输入：

x[np.argpartition(x, -3)[-3:]]

输出：

array([2, 4, 3])

输入：

x[np.argpartition(x, -2)[-2:]] 

输出：

array([3, 4])

输入：

x[np.argpartition(x, -1)[-1:]] 

输出：

array([4])

如上代码所示：
若想取从小到大排序前n的数就kth=-n，并取[-n:]索引就可以得到最大的n个数。