Alexis Métaireau Alexis Métaireau - 2 years ago 235
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 Source

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.

Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download