Question

我在StackOverflow上经历了几个问题但未找到相关答案。我想从numpy ndarray获取k个最大值的索引。 This link讨论了相同的但是对于一维数组。对于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

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

1 个答案: