CODEWITHSUNDEEP
pythonsortingnumpy
gerry
gerry

Reputation: 1569

How to get the index of the sorted list for the original list in Python?

The argsort function of numpy returns the original index of the sorted list. Now I want the index of the sorted list for the original list. Is there a function or an elegant way to do this ?

For example:

>>> import numpy as np
>>> a = np.array([2, 8 , 5, 6])
>>> np.argsort(a)
array([0, 2, 3, 1])
>>> function(a)
array([0, 3, 1, 2])

Upvotes: 6

Views: 4017

Answers (2)

Jaime
Jaime

Reputation: 67427

While the double argsort trick works, it is not very efficient. You can get better performance by using fancy indexing:

>>> argsort = a.argsort()
>>> rev_argsort = np.empty(argsort.shape, dtype=np.intp)
>>> rev_argsort[argsort] = np.arange(len(a))
>>> rev_argsort
array([0, 3, 1, 2])

It is more verbose, but it has linear complexity instead of the linearithmic complexity of argsort. In practice this means that, for sufficiently large arrays, the above code will run twice as fast as a double argsort, as the time to create and fill rev_argsort will be negligible against that of the first argsort.

Upvotes: 7

Alex Riley
Alex Riley

Reputation: 176750

You could always call argsort twice:

>>> a.argsort().argsort()
array([0, 3, 1, 2])

As far as I know, there's no "double-argsort" function available in NumPy, but applying argsort twice to an array is a common way to calculate the ranks of the values (cf. here).

Upvotes: 9

Related Questions