Alexis Métaireau - 1 year ago 140

Python Question

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

`function(array, n=3)`

`[4, 3, 1]`

Thanks :)

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.

Source (Stackoverflow)