numpy.argpartition() 函数介绍

---

官方介绍

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

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.

---

我的理解

该函数是用来取一个数组的前k大或前k小的数的。那么为什么不用sort要用numpy.argpartition() 呢？因为numpy.argpartition() 只是将比第k个位置的数小的都放到它前面，比第k个数大的都放到它后面，至于这些数内部的排序，numpy.argpartition() 并不保证，因此在算法的复杂度上是O(n)。

当我们不需要对整个数组进行全排列而是只想要比如前100大的数，那么numpy.argpartition() 是一个性能上更优的选择（当然，我觉得用sort在大部分情况下性能的差距并不大）

---

具体使用

创建一个数组

x = np.array([11, 4, 7, 1, 6, 9, 5, 8, 12, 5])

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

其输出为：

array([ 4,  5,  5,  1,  6,  7,  8,  9, 12, 11])
array([ 4,  5,  5,  1,  6,  7,  8,  9, 12, 11])
array([ 4,  5,  5,  1,  6,  7,  8,  9, 12, 11])
array([ 1,  4,  5,  5,  6,  9,  7,  8, 12, 11])

我们具体使用起来是这样的：

# 取前5大的数
x[np.argpartition(x, -5)][-5:]
# 取前3大的数
x[np.argpartition(x, -3)][-3:]
# 取前5小的数
x[np.argpartition(x, 5)][:5]
# 取前3小的数
x[np.argpartition(x, 3)][:3]

其输出为：

array([ 7,  8,  9, 12, 11])
array([ 9, 12, 11])
array([4, 5, 5, 1, 6])
array([1, 4, 5])

---

总结

以上