Alexis Métaireau Alexis Métaireau - 7 months ago 71
Python Question

How to get indices of N maximum values in a numpy array?

Numpy proposes a way to get the index of the maximum value of an array via

np.argmax
.

I would like a similar thing, but returning the indexes of the N maximum values.

For instance, if I have an array
[1, 3, 2, 4, 5]
, it
function(array, n=3)
would return
[4, 3, 1]
.

Thanks :)

NPE NPE
Answer

The simplest I've been able to come up with is:

In [1]: import numpy as np

In [2]: arr = np.array([1, 3, 2, 4, 5])

In [3]: arr.argsort()[-3:][::-1]
Out[3]: array([4, 3, 1])

This involves a complete sort of the array. I wonder if numpy provides a built-in way to do a partial sort; so far I haven't been able to find one.

If this solution turns out to be too slow (especially for small n), it may be worth looking at coding something up in Cython.