Sort a list then give the indexes of the elements in their original order

Question

I have an array of n numbers, say [1,4,6,2,3]. The sorted array is [1,2,3,4,6], and the indexes of these numbers in the old array are 0, 3, 4, 1, and 2. What is the best way, given an array of n numbers, to find this array of indexes?

My idea is to run order statistics for each element. However, since I have to rewrite this function many times (in contest), I'm wondering if there's a short way to do this.

Solution 1

>>> a = [1,4,6,2,3]
>>> [b[0] for b in sorted(enumerate(a),key=lambda i:i[1])]
[0, 3, 4, 1, 2]

Explanation:

enumerate(a) returns an enumeration over tuples consisting of the indexes and values in the original list: [(0, 1), (1, 4), (2, 6), (3, 2), (4, 3)]

Then sorted with a key of lambda i:i[1] sorts based on the original values (item 1 of each tuple).

Finally, the list comprehension [b[0] for b in ...] returns the original indexes (item 0 of each tuple).

Solution 2

Using numpy arrays instead of lists may be beneficial if you are doing a lot of statistics on the data. If you choose to do so, this would work:

import numpy as np
a = np.array( [1,4,6,2,3] )
b = np.argsort( a )

argsort() can operate on lists as well, but I believe that in this case it simply copies the data into an array first.

Solution 3

Here is another way:

>>> sorted(xrange(len(a)), key=lambda ix: a[ix])
[0, 3, 4, 1, 2]

This approach sorts not the original list, but its indices (created with xrange), using the original list as the sort keys.