从numpy数组中获取最大或最小n元素？(最好不要压扁)

Question

我知道我可以获得最小值或最大值:

max(matrix)
min(matrix)

从一个numpy矩阵/向量.这些值的索引由以下方式返回:

argmax(matrix)
argmin(matrix)

所以例如当我有一个5x5矩阵时:

a = np.arange(5*5).reshape(5, 5) + 10

# array([[10, 11, 12, 13, 14],
#        [15, 16, 17, 18, 19],
#        [20, 21, 22, 23, 24],
#        [25, 26, 27, 28, 29],
#        [30, 31, 32, 33, 34]])

我可以通过以下方式获得最大值:

In [86]: np.max(a) # getting the max-value out of a
Out[86]: 34

In [87]: np.argmax(a) # index of max-value 34 is 24 if array a were flattened
Out[87]: 24

...但是获得最大或最小n元素的最有效方法是什么？

所以我们说出来的一个我想有5最高和最低5元.这应该[30, 31, 32, 33, 34]分别[20, 21, 22, 23, 24]为我们的指数返回5个最高值.同样[10, 11, 12, 13, 14],对于5个最低值和[0, 1, 2, 3, 4]5个最低元素的索引.

什么是一个有效,合理的解决方案？

我的第一个想法是对数组进行展平和排序,并采用最后和前5个值.然后,我在原始2D矩阵中搜索这些值的索引.虽然这个程序工作扁平化+排序效率不高......有没有人知道更快的解决方案？

另外我想要原始2D数组的索引而不是展平的索引.因此24,np.argmax(a)我希望不会返回(4, 4).

Answer 1

获取数组中最大或最小值索引的标准方法是使用np.argpartition. 该函数使用 introselect 算法并以线性复杂度运行 - 这比对较大数组进行完全排序（通常为 O(n log n)）要好。

默认情况下，该函数沿数组的最后一个轴工作。要考虑整个数组，您需要使用ravel(). 例如，这是一个随机数组a：

>>> a = np.random.randint(0, 100, size=(5, 5))
>>> a
array([[60, 68, 86, 66,  9],
       [66, 26, 83, 87, 50],
       [41, 26,  0, 55,  9],
       [57, 80, 71, 50, 22],
       [94, 30, 95, 99, 76]])

然后，要获取（展平的）二维数组中五个最大值的索引，请使用：

>>> i = np.argpartition(a.ravel(), -5)[-5:] # argpartition(a.ravel(), 5)[:5] for smallest
>>> i
array([ 2,  8, 22, 23, 20])

要获取 中这些位置的相应二维索引a，请使用unravel_index：

>>> i2d = np.unravel_index(i, a.shape)
>>> i2d
(array([0, 1, 4, 4, 4]), array([2, 3, 2, 3, 0]))

然后索引a返回i2d五个最大值：

>>> a[i2d]
array([86, 87, 95, 99, 94])