python - 从 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 个元素的最有效方法是什么？

所以让我们说出 a 我想要 5 个最高和 5 个最低的元素。这应该为我返回 [30, 31, 32, 33, 34] 5 个最高值分别为 [20, 21, 22, 23, 24] 作为它们的索引。同样，[10, 11, 12, 13, 14] 用于 5 个最低值，[0, 1, 2, 3, 4] 用于 5 个最低值的索引元素。

什么是有效、合理的解决方案？

我的第一个想法是展平和排序数组并取最后和前 5 个值。之后，我在原始二维矩阵中搜索这些值的索引。 虽然此过程可以展平 + 排序但效率不高...有人知道更快的解决方案吗？

此外，我想要原始二维数组的索引，而不是展平数组的索引。所以不是 np.argmax(a) 返回的 24 我想要 (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]))

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

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