[Numpy-discussion] Find indices of largest elements

[Keith Goodman]

On Wed, Apr 14, 2010 at 8:16 AM, Nikolaus Rath <Nikolaus at rath.org> wrote:
> Hello,
>
> How do I best find out the indices of the largest x elements in an
> array?
>
> Example:
>
> a = [ [1,8,2], [2,1,3] ]
> magic_function(a, 2) == [ (0,1), (1,2) ]
>
> Since the largest 2 elements are at positions (0,1) and (1,2).

Here's a quick way to rank the data if there are no ties and no NaNs:

>> shape = (3,2)
>> x = np.random.rand(*shape)
>> x
array([[ 0.83424288,  0.17821326],
       [ 0.62718311,  0.63514286],
       [ 0.18373934,  0.90634162]])
>> r = x.reshape(-1).argsort().argsort().reshape(shape)
>> r
array([[4, 0],
       [2, 3],
       [1, 5]])

To find the indices you can use where:

>> r < 2
array([[False,  True],
       [False, False],
       [ True, False]], dtype=bool)
>> np.where(r < 2)
   (array([0, 2]), array([1, 0]))

...but the indices will not be in order.