python - Get the indices of N highest values in an ndarray

Question

Considering an histogram of shape 100x100x100, I would like to find the 2 highest values a and b, and their indices (a1, a2, a3) and (b1, b2, b3), such as:

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

We can easily get the highest value with hist.max(), but how can we get the X highest values in a ndarray?

I understand that one normally uses np.argmax to retrieve the value indices, but in that case:

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

How can I get a shape (3), a tuple with one value per dimension?

Answer 1

You can use numpy.argpartition on flattened version of array first to get the indices of top k items, and then you can convert those 1D indices as per the array's shape using numpy.unravel_index:

>>> 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