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

Question

我在StackOverflow上经历了几个问题，但找不到相关的答案。 我想从numpy ndarray获取k个最大值的索引。 此链接讨论相同但对于1D阵列。 用于2D数组的np.argsort导致逐行排序元素。 即

Note: array elements are not unique.

输入：

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)

也，

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)

但我在这里丢失了数组结构，无法兑换元素的原始索引。

任何帮助都是值得赞赏的。

Answer 1

np.argpartition和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和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