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


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

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