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

Question

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

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