[python] How to get indices of a sorted array in Python

Essentially you need to do an argsort, what implementation you need depends if you want to use external libraries (e.g. NumPy) or if you want to stay pure-Python without dependencies.

The question you need to ask yourself is: Do you want the

• indices that would sort the array/list
• indices that the elements would have in the sorted array/list

Unfortunately the example in the question doesn't make it clear what is desired because both will give the same result:

>>> arr = np.array([1, 2, 3, 100, 5])

>>> np.argsort(np.argsort(arr))
array([0, 1, 2, 4, 3], dtype=int64)

>>> np.argsort(arr)
array([0, 1, 2, 4, 3], dtype=int64)

Choosing the argsort implementation

If you have NumPy at your disposal you can simply use the function numpy.argsort or method numpy.ndarray.argsort.

An implementation without NumPy was mentioned in some other answers already, so I'll just recap the fastest solution according to the benchmark answer here

def argsort(l):
    return sorted(range(len(l)), key=l.__getitem__)

Getting the indices that would sort the array/list

To get the indices that would sort the array/list you can simply call argsort on the array or list. I'm using the NumPy versions here but the Python implementation should give the same results

>>> arr = np.array([3, 1, 2, 4])
>>> np.argsort(arr)
array([1, 2, 0, 3], dtype=int64)

The result contains the indices that are needed to get the sorted array.

Since the sorted array would be [1, 2, 3, 4] the argsorted array contains the indices of these elements in the original.

• The smallest value is 1 and it is at index 1 in the original so the first element of the result is 1.
• The 2 is at index 2 in the original so the second element of the result is 2.
• The 3 is at index 0 in the original so the third element of the result is 0.
• The largest value 4 and it is at index 3 in the original so the last element of the result is 3.

Getting the indices that the elements would have in the sorted array/list

In this case you would need to apply argsort twice:

>>> arr = np.array([3, 1, 2, 4])
>>> np.argsort(np.argsort(arr))
array([2, 0, 1, 3], dtype=int64)

In this case :

• the first element of the original is 3, which is the third largest value so it would have index 2 in the sorted array/list so the first element is 2.
• the second element of the original is 1, which is the smallest value so it would have index 0 in the sorted array/list so the second element is 0.
• the third element of the original is 2, which is the second-smallest value so it would have index 1 in the sorted array/list so the third element is 1.
• the fourth element of the original is 4 which is the largest value so it would have index 3 in the sorted array/list so the last element is 3.