numpy.argsort() in Python

The NumPy module provides a function argsort(), returns the indices which would sort an array.

The NumPy module provides a function for performing an indirect sort along with the given axis with the help of the algorithm specified by the keyword. This function returns an array of indices of the same shape as 'a', which would sort the array.

Syntax

numpy.argsort(a, axis=-1, kind=None, order=None)

Parameters

These are the following parameters in numpy.argsort() function:

a: array_like

This parameter defines the source array which we want to sort.

axis: int or None(optional)

This parameter defines the axis along which the sorting is performed. By default, the axis is -1. If we set this parameter to None, the flattened array is used.

kind: {'quicksort','mergesort','heapsort','stable'}(optional)

This parameter defines the sorting algorithm. By default, the algorithm is quicksort. Both mergesort and stable are using time sort under the covers. The actual implementation will vary with the data type. The mergesort option is retained for backward compatibility.

order: str or list of str(optional)

If 'a' is an array with defined fields, this argument specifies which fields to compare first, second, etc. The single field can be specified as a string, and not all fields need to be specified. But unspecified fields will still use, in the order in which they come up in the dtype, to break the ties.

Returns: index_array: ndarray, int

This function returns an array of indices which sort 'a' along with the specified axis. If 'a' is 1-D, a[index_array] yields a sorted 'a'. More generally, np.take_along_axis(arr1, index_array, axis=axis) always yields the sorted 'a', irrespective of dimensionality.

Example 1: np.argsort()

import numpy as np
a=np.array([456,11,63])
a
b=np.argsort(a)
b

In the above code

• We have imported numpy with alias name np.
• We have created an array 'a' using np.array() function.
• We have declared the variable 'b' and assigned the returned value of np.argsort() function.
• We have passed the array 'a' in the function.
• Lastly, we tried to print the value of b.

In the output, a ndarray has been shown that contains the indices(indicate the position of the element for the sorted array) and dtype.

Output:

Example 2: For 2-D array( sorts along first axis (down))

import numpy as np
a = np.array([[0, 5], [3, 2]])
indices = np.argsort(a, axis=0)  
indices

Output:

Example 3: For 2-D array(alternative of axis=0)

import numpy as np
a = np.array([[0, 5], [3, 2]])
indices = np.argsort(a, axis=0)
indices
np.take_along_axis(a, indices, axis=0)

In the above code

• We have imported numpy with alias name np.
• We have created a 2-D array 'a' using np.array() function.
• We have declared variable indices and assigned the returned value of np.argsort() function.
• We have passed the 2-D array 'a' and axis as 0.
• Next, we used the take_along_axis() function and passed the source array, indices, and axis.
• This function has returned the sorted 2-D array.

In the output, a 2-D array with sorted elements has been shown.

Output:

Example 4: For 2-D array( sorts along last axis (across))

import numpy as np
a = np.array([[0, 5], [3, 2]])
indices = np.argsort(a, axis=1)  
indices

Output:

Example 5: For 2-D array(alternative of axis=1)

import numpy as np
a = np.array([[0, 5], [3, 2]])
indices = np.argsort(a, axis=1)
indices
np.take_along_axis(a, indices, axis=1)

Output:

Example 6: For N-D array

import numpy as np
a = np.array([[0, 5], [3, 2]])
indices = np.unravel_index(np.argsort(a, axis=None), a.shape)
indices
a[indices]  # same as np.sort(a, axis=None)

Output:

In the above code

• We have imported numpy with alias name np.
• We have created a 2-D array 'a' using np.array() function.
• We have declared a variable 'indices' and assigned the returned value of np.unravel_index() function.
• We have passed the np.argsort() function and shape of the array 'a'.
• We have passed the 2-D array 'a' and axis as 1 in argsort() function.
• Next, we tried to print the value of indices and a[indices].

In the output, an N-D array with sorted elements has been shown.

Example 7: Sorting with keys

import numpy as np
a= np.array([(0, 5), (3, 2)], dtype=[('x', '<i4'), ('y', '<i4')])
a
b=np.argsort(a, order=('x','y'))
b
c=np.argsort(a, order=('y','x'))
c

Output:

In the above code

• We have imported numpy with alias name np.
• We have created a 2-D array 'a' using np.array() function with dtype=[('x', '<i4'), ('y', '<i4')].
• We have declared the variables 'b' and 'c' and assigned the returned value of np.argsort() function.
• We have passed the array 'a' and order as an argument in the function.
• Lastly, we tried to print the value of 'b' and 'c'.

In the output, a sorted array has been shown with dtype=[('x', '<i4'), ('y', '<i4')]