如何從一個numpy多維數組中獲取k個最大值的索引

Question

我在StackOverflow上經歷了幾個問題，但找不到相關答案。我想從numpyndarray中獲取k​​個最大值的索引。 This link討論相同，但爲一維數組。針對二維數組的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)) 

談兩個版本k_largest_index_argpartition_v1和k_largest_index_argpartition_v2之間argpartition

不同的是，我們是如何使用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