获取ndarray中N个最高值的索引

Question

考虑到形状100x100x100的直方图,我想找到2个最高值a和b,以及它们的索引(a1,a2,a3)和(b1,b2,b3),例如:

hist[a1][a2][a3] = a
hist[b1][b2][b3] = b

我们可以使用hist.max()轻松获得最高值,但是如何在ndarray中获得X最高值？

我知道通常使用np.argmax来检索值索引,但在这种情况下:

hist.argmax().shape = ()  # single value
for i in range(3):
    hist.argmax(i).shape = (100, 100)

我怎样才能得到一个形状(3),一个每个维度有一个值的元组？

Answer 1

您可以numpy.argpartition首先在平顶版本的数组上使用以获取顶部k项目的索引,然后您可以使用numpy.unravel_index以下内容根据数组的形状转换这些1D索引:

>>> arr = np.arange(100*100*100).reshape(100, 100, 100)
>>> np.random.shuffle(arr)
>>> indices =  np.argpartition(arr.flatten(), -2)[-2:]
>>> np.vstack(np.unravel_index(indices, arr.shape)).T
array([[97, 99, 98],
       [97, 99, 99]])
)
>>> arr[97][99][98]
999998
>>> arr[97][99][99]
999999