如何从 numpy 多维数组中获取 k 个最大值的索引

我问了几个问题StackOverflow但找不到相关答案。我想从 a 中获取 k 个最大值的索引numpy ndarray. 这个链接讨论相同但针对一维数组。np.argsort对于二维数组导致元素按行排序。 IE

Note: array elements are not unique.

input:

import numpy as np
n = np.arange(9).reshape(3,3)
>>> n
array([[0, 1, 2],
   [3, 4, 5],
   [6, 7, 8]])
s = n.argsort()
>>> s
array([[0, 1, 2],
   [0, 1, 2],
   [0, 1, 2]], dtype=int32)

Also,

import numpy as np
n = np.arange(9).reshape(3,3)
s = n.argsort(axis=None)
>>>s
array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=int32)

但我在这里失去了数组结构，无法恢复元素的原始索引。

任何形式的帮助表示赞赏。

几种方法np.argpartition and np.argsort对于 ndarrays -

def k_largest_index_argpartition_v1(a, k):
    idx = np.argpartition(-a.ravel(),k)[:k]
    return np.column_stack(np.unravel_index(idx, a.shape))

def k_largest_index_argpartition_v2(a, k):
    idx = np.argpartition(a.ravel(),a.size-k)[-k:]
    return np.column_stack(np.unravel_index(idx, a.shape))

def k_largest_index_argsort(a, k):
    idx = np.argsort(a.ravel())[:-k-1:-1]
    return np.column_stack(np.unravel_index(idx, a.shape))

关于两个版本的讨论argpartition

之间的区别k_largest_index_argpartition_v1 and k_largest_index_argpartition_v2就是我们使用的方式argparition。在第一个版本中，我们对输入数组取反，然后使用argpartition获得最小的索引k指数，从而有效地获得最大的k索引，而在第二个版本中，我们得到第一个a.size-k最小的索引，然后我们选择剩下的最大的k指数。

另外，这里值得一提的是argpartition，我们没有按排序顺序获取索引。如果需要排序顺序，我们需要输入范围数组np.argpartition，正如本文中提到的post.

样本运行 -

1) 2D 情况:

In [42]: a    # 2D array
Out[42]: 
array([[38, 14, 81, 50],
       [17, 65, 60, 24],
       [64, 73, 25, 95]])

In [43]: k_largest_index_argsort(a, k=2)
Out[43]: 
array([[2, 3],
       [0, 2]])

In [44]: k_largest_index_argsort(a, k=4)
Out[44]: 
array([[2, 3],
       [0, 2],
       [2, 1],
       [1, 1]])

In [66]: k_largest_index_argpartition_v1(a, k=4)
Out[66]: 
array([[2, 1], # Notice the order is different
       [2, 3],
       [0, 2],
       [1, 1]])

2）3D案例：

In [46]: a # 3D array
Out[46]: 
array([[[20, 98, 27, 73],
        [33, 78, 48, 59],
        [28, 91, 64, 70]],

       [[47, 34, 51, 19],
        [73, 38, 63, 94],
        [95, 25, 93, 64]]])

In [47]: k_largest_index_argsort(a, k=2)
Out[47]: 
array([[0, 0, 1],
       [1, 2, 0]])

运行时测试 -

In [56]: a = np.random.randint(0,99999999999999,(3000,4000))

In [57]: %timeit k_largest_index_argsort(a, k=10)
1 loops, best of 3: 2.18 s per loop

In [58]: %timeit k_largest_index_argpartition_v1(a, k=10)
10 loops, best of 3: 178 ms per loop

In [59]: %timeit k_largest_index_argpartition_v2(a, k=10)
10 loops, best of 3: 128 ms per loop