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

Answer 2

You could use where:

a=np.random.random((100,100,100))
np.where(a==a.max())
(array([46]), array([62]), array([61]))

to get in a single array:

np.hstack(np.where(a==a.max()))
array([46, 62, 61])

and, as the OP asked for a tuple:

tuple(np.hstack(np.where(a==a.max())))
(46, 62, 61)

EDIT:

To get the indices of the N largest sets you could use the nlargest function from the heapq module:

N=3
np.where(a>=heapq.nlargest(3,a.flatten())[-1])
(array([46, 62, 61]), array([95, 85, 97]), array([70, 35,  2]))