How to Get the Indices of the Sorted Array using NumPy in Python

The numpy.argsort() function is used to get indices that would sort a NumPy array. Instead of returning sorted array, it returns positions of the elements in order that sorts the array. This is helpful when you need both the sorted order and the original indices.

import numpy as np
arr = np.array([30, 10, 20])
idx = np.argsort(arr)
print(idx)

[1 2 0]

Explanation: Element 10 is the smallest -> index 1 comes first. Next is 20 -> index 2. Last is 30 -> index 0. So, sorted order of indices is [1, 2, 0].

Syntax

numpy.argsort(arr, axis=-1, kind='quicksort', order=None)

Parameters:

• arr: Input array to be sorted.
• axis: Axis along which to sort. Default is -1 (last axis).
• kind: Sorting algorithm to use: 'quicksort', 'mergesort', 'heapsort' or 'stable'. Default is 'quicksort'.
• order: When sorting structured arrays, this specifies which field to sort by.

Return Value: Returns an array of indices of the same shape as arr. These indices can be used to reconstruct sorted array.

Examples

Example 1: In this example, we find the indices that would sort a one-dimensional array.

import numpy as np

arr = np.array([10, 52, 62, 16, 16, 54, 453])
print("Array:", arr)

idx = np.argsort(arr)
print("Sorted Indices:", idx)

Array: [ 10  52  62  16  16  54 453]
Sorted Indices: [0 3 4 1 5 2 6]

Example 2: This example shows that if the array is already sorted, indices are returned in sequential order.

import numpy as np

arr = np.array([1, 2, 3, 4, 5])
print("Array:", arr)

idx = np.argsort(arr)
print("Sorted Indices:", idx)

Array: [1 2 3 4 5]
Sorted Indices: [0 1 2 3 4]

Example 3: Here we demonstrate how argsort() works along rows and columns of a 2-D array using different sorting algorithms.

import numpy as np 

arr1 = np.array([[2, 0, 1], [5, 4, 3]]) 
print("Input Array:") 
print(arr1)  

arr2 = np.argsort(arr1, kind='mergesort', axis=0) 
print("Sorted Indices along Axis 0:")
print(arr2) 

arr3 = np.argsort(arr1, kind='heapsort', axis=1) 
print("Sorted Indices along Axis 1:")
print(arr3) 

Input Array:
[[2 0 1]
 [5 4 3]]
Sorted Indices along Axis 0:
[[0 0 0]
 [1 1 1]]
Sorted Indices along Axis 1:
[[1 2 0]
 [2 1 0]]