How do I get indices of N maximum values in a NumPy array?

后端 未结 18 1478
长情又很酷
长情又很酷 2020-11-22 04:25

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

18条回答
  •  借酒劲吻你
    2020-11-22 05:09

    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.

提交回复
热议问题