np.argsort() — A Simpe Illustrated Guide

In Python, the numpy.argsort() function returns the indices that would sort an array in ascending order. 

Here is the argument table of the numpy.argsort() function.

If it sounds great to you, please continue reading, and you will fully understand the numpy.argsort() function through Python code snippets and vivid visualization.

This tutorial is about numpy.argsort() function. 

  • Concretely, I will introduce its syntax and arguments. 
  • Then, you will learn some basic examples of this function.
  • Finally, I will address three top questions about numpy.argsort(), including np.argsort descending, np.argsort 2d array/axis, and np.argsort order.

You can find all codes in this tutorial here.

Syntax and Arguments

Here is the syntax of numpy.argsort():

# syntax
numpy.argsort(a[, axis=-1[, kind=None[, order=None]]])

Here is the argument table of numpy.argsort():

ArgumentAcceptDescription
aarray_likeArray to sort.
axisint or None, optionalAxis along which to sort. The default is -1 (the last axis). If None, the flattened array is used.
For more information about a flattened array, please refer to numpy.ndarray.flatten() function.
kind{'quicksort', 'mergesort', 'heapsort', 'stable'}, optionalSorting algorithm. The default is 'quicksort'.
Note that both โ€˜stableโ€™ and โ€˜mergesortโ€™ use timsort under the covers and, in general, the actual implementation will vary with data type.
The โ€˜mergesortโ€™ option is retained for backwards compatibility.
orderstr or list of str, optionalThis argument specifies the order in which to compare field.

Generally, you only need to deploy the a and axis arguments. And I will clearly explain the axis and order argument later.

The output of numpy.argsort() function is an array of indices which sort the input array along with the specified axis. 

On top of that, If the input array 'a' is 1-D, a[output] yields a sorted 'a'. More generally, np.take_along_axis(a, output, axis=axis) always yields the sorted 'a', irrespective of dimensionality. We will see more examples later in this article. ๐Ÿ™‚

Basic Examples

Here is a one-dimensional array code example:

import numpy as np

one_dim = np.array([1, 5, 4, 0, 3])
sort_index = np.argsort(one_dim)  # axis defaults to -1 (the last axis)
print(f'Unsorted array: {one_dim}')
print(f'Sorted array index: {sort_index}')
print(f'Sorted array: {one_dim[sort_index]}')

Output:

In this basic example, we use the numpy.argsort() function to return the index of the sorted input array (in an ascending order) and then index the sorted array using the output.

np.argsort() descending

We can also return the index that sorts the input array in a descending order using the [::-1] reverse trick.

Here is a one-dimensional array code example:

import numpy as np

one_dim = np.array([1, 2, 3, 4, 5])
# use the [::-1] to reverse the ascending order to descending order.
sort_index = np.argsort(one_dim)[::-1]
print(f'Unsorted array: {one_dim}')
print(f'Sorted array index: {sort_index}')
print(f'Sorted array: {one_dim[sort_index]}')

Output:

Yes, it is just like reversing a string. We can simply add the [::-1] to the output of the np.argsort() function to get a descending order sorted index.

np.argsort() 2d array / axis

So far, weโ€™ve only seen some 1d array examples.

In this part, I will show you how to deploy the axis argument with some 2d array examples!

By the way, you can always use the np.take_along_axis(input_array, sort_index, axis=axis) to get the sorted 'a', irrespective of dimensionality. 

Here is the argument table for reference:

Here is the 2d array example with axis=0:

import numpy as np

# Here is the 2d array example with axis=0:
two_dim = np.array([[1, 2, 3], [3, 2, 1]])
sort_index = np.argsort(two_dim, axis=0)

print(f"Unsorted array: {two_dim}")
print(f"Sorted array index: {sort_index}")
print(f"Sorted array: {np.take_along_axis(two_dim, sort_index, axis=0)}")

Output:

Here is the 2d array example with axis=1:

# Here is the 2d array example with axis=1:
import numpy as np

two_dim = np.array([[1, 2, 3], [3, 2, 1]])
sort_index = np.argsort(two_dim, axis=1)

print(f"Unsorted array: {two_dim}")
print(f"Sorted array index: {sort_index}")
print(f"Sorted array: {np.take_along_axis(two_dim, sort_index, axis=1)}")

Output:

Here is the 2d array example with axis=None:

# Here is the 2d array example with axis=None:
import numpy as np

two_dim = np.array([[1, 2, 3], [3, 2, 1]])
sort_index = np.argsort(two_dim, axis=None)

print(f"Unsorted array: {two_dim}")
print(f"Sorted array index: {sort_index}")
print(f"Sorted array: {np.take_along_axis(two_dim, sort_index, axis=None)}")

Output:

np.argsort() order

Although I said that you would probably only need to deploy the a and axis argument, I think the order argument is probably confusing to you. So I will give you an example to help you understand what it means!

In the official documentation, the order argument goes,


โ€œWhen a is an array with fields defined, this argument specifies which fields to compare first, second, etc. A single field can be specified as a string, and not all fields need be specified, but unspecified fields will still be used, in the order in which they come up in the dtype, to break ties.โ€


Intuitively, the order argument sets the sorting order for an input array with defined fields. You will have a better sense of what it means after seeing the next code examples.

Here are the 1d array examples for order argument:

import numpy as np

# order = x -> y
one_dim = np.array([(1, 2), (1, 1), (2, 2), (2, 1)],
                   dtype=np.dtype([('x', int), ('y', int)]))
sort_index = np.argsort(one_dim)  # or np.argsort(x, order=('x', 'y'))
print(f'Unsorted array: {one_dim}')
print(f'Sorted array index: {sort_index}')
print(f'Sorted array: {one_dim[sort_index]}')
print('-' * 85)


# order = y -> x
one_dim = np.array([(1, 2), (1, 1), (2, 2), (2, 1)],
                   dtype=np.dtype([('x', int), ('y', int)]))
sort_index = np.argsort(one_dim, order=('y', 'x'))
print(f'Unsorted array: {one_dim}')
print(f'Sorted array index: {sort_index}')
print(f'Sorted array: {one_dim[sort_index]}')

Output:

Summary

Thatโ€™s it for our np.argsort() article. 

We learned about its syntax, arguments, and basic examples. 

We also worked on the top three questions about the np.argsort() function, ranging from np.argsort() descending, np.argsort() 2d array/axis, and np.argsort() order

Actually, Python provides quite a lot sorting related functions aside from the numpy.argsort() function. 

  • If you want to have more flexibility with ties situation, check out the scipy.stats.rankdata() function. 
  • If you want to indirectly sort an array with multiple keys, check out the numpy.lexsort() function. 
  • If you want to directly sort an array, check out the numpy.sort() and numpy.ndarray.sort() function. 

Of course, if you want me to explain more hard-to-understand function, just let me know. ๐Ÿ™‚

Hope you enjoy all this and happy coding!