如何獲得多維numpy數組中最大n值的索引

Question

我想獲得多維numpy數組中最大n值的索引。爲了得到一維numpy數組中最大n值的索引，我找到了this。在python的交互式shell中測試之後，似乎bottleneck.argpartsort不能對多維numpy數組產生影響。爲了獲得多維numpy數組中最大值的索引，我找到了this。它不能得到最大的n。我可以給出的方法是將多維numpy數組轉換爲{value:index}（由元組存在的索引）列表，然後按值對列表進行排序，並獲取索引。有什麼更容易或更多的表現？

Answer 1

我沒有獲得bottleneck，所以在這個例子中，我使用argsort，但你應該能夠以同樣的方式使用它：

#!/usr/bin/env python 
import numpy as np 
N = 4 
a = np.random.random(20).reshape(4, 5) 
print(a) 

# Convert it into a 1D array 
a_1d = a.flatten() 

# Find the indices in the 1D array 
idx_1d = a_1d.argsort()[-N:] 

# convert the idx_1d back into indices arrays for each dimension 
x_idx, y_idx = np.unravel_index(idx_1d, a.shape) 

# Check that we got the largest values. 
for x, y, in zip(x_idx, y_idx): 
    print(a[x][y])