[Numpy-discussion] Find indices of largest elements

[Keith Goodman]

On Wed, Apr 14, 2010 at 8:49 AM, Keith Goodman <kwgoodman at gmail.com> wrote:
> 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.

...or if you need the indices in order:

>> shape = (3,2)
>> x = np.random.rand(*shape)
>> x
array([[ 0.52420123,  0.43231286],
       [ 0.97995333,  0.87416228],
       [ 0.71604075,  0.66018382]])
>> r = x.reshape(-1).argsort().argsort()
>> np.unravel_index(r[0], shape)
   (0, 1)
>> np.unravel_index(r[1], shape)
   (0, 0)
>> np.unravel_index(r[2], shape)
   (2, 1)
>> np.unravel_index(r[-1], shape)
   (1, 0)