numpy.argsort() in Python

numpy.argsort() is a function in NumPy that returns the indices that would sort an array. In other words, it gives you the indices that you would use to reorder the elements in an array to be in sorted order. Example:

import numpy as np

a = np.array([2, 0, 1, 5, 4, 1, 9])
idx = np.argsort(a)

print("Original array:", a)
print("Indices to sort:", idx)
print("Sorted array:", a[idx])

Original array: [2 0 1 5 4 1 9]
Indices to sort: [1 2 5 0 4 3 6]
Sorted array: [0 1 1 2 4 5 9]

Explanation: np.argsort(a) returns [1 2 5 0 4 3 6], the indices that would sort a. and a[idx] gives the sorted array.

Syntax

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

Parameters:

• arr : [array_like] Input array. 
• axis : [int or None] Axis along which to sort. If None, the array is flattened before sorting. The default is -1, which sorts along the last axis. 
• kind : [‘quicksort’, ‘mergesort’, ‘heapsort’]Selection algorithm. Default is ‘quicksort’. 
• order : [str or list of str] When arr is an array with fields defined, this argument specifies which fields to compare first, second, etc. 

Return: [index_array, ndarray] Array of indices that sort arr along the specified axis. If arr is one-dimensional then arr[index_array] returns a sorted arr.

Examples of numpy.argsort()

Example 1: This example demonstrates how to use numpy.argsort() with different axes to return indices that would sort a 2D NumPy array along rows and columns.

import numpy as np

a = np.array([[2, 0, 1], [5, 4, 3]])
print("Axis 0:\n", np.argsort(a, axis=0))
print("Axis 1:\n", np.argsort(a, axis=1))

Axis 0:
 [[0 0 0]
 [1 1 1]]
Axis 1:
 [[1 2 0]
 [2 1 0]]

Explanation:

• Axis 0 (columns): The values in each column are compared top-to-bottom.
• Axis 1 (rows): The values in each row are compared left-to-right.

Example 2: This example finds the indices of the two largest elements in a NumPy array and retrieves their values in descending order.

import numpy as np

x=np.array([12,43,2,100,54,5,68])

print(np.argsort(x))
print(np.argsort(x)[-2:])
print(np.argsort(x)[-2:][::-1])
print(x[np.argsort(x)[-2:][::-1]])

[2 5 0 1 4 6 3]
[6 3]
[3 6]
[100  68]

Explanation:

• np.argsort(x) returns the indices that would sort the array.
• np.argsort(x)[-2:] selects the indices of the two largest elements.
• Reversing the indices with [::-1] places the largest element first.
• Finally, the array is indexed with the reversed indices to extract the two largest elements.

Example 3: This example assigns ranks to array elements based on their values, with the smallest value ranked 0, the next smallest 1 and so on.

import numpy as np
a = np.array([50, 10, 20, 30])

res = np.argsort(np.argsort(a))
print(res)

[3 0 1 2]

Explanation:

• First argsort gives the sort order [1 2 3 0]
• Second argsort gives rank for each item: where it falls in the sorted order.