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

后端 未结 18 1428
长情又很酷
长情又很酷 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:10

    This will be faster than a full sort depending on the size of your original array and the size of your selection:

    >>> A = np.random.randint(0,10,10)
    >>> A
    array([5, 1, 5, 5, 2, 3, 2, 4, 1, 0])
    >>> B = np.zeros(3, int)
    >>> for i in xrange(3):
    ...     idx = np.argmax(A)
    ...     B[i]=idx; A[idx]=0 #something smaller than A.min()
    ...     
    >>> B
    array([0, 2, 3])
    

    It, of course, involves tampering with your original array. Which you could fix (if needed) by making a copy or replacing back the original values. ...whichever is cheaper for your use case.

提交回复
热议问题