NumPy - Sorting, Searching and Counting

The NumPy package contains a number of functions which provides all the functionality required for sorting, searching and counting. Below mentioned are the most frequently used functions to perform such operations on an array.

Lets discuss these functions in detail:

numpy.sort() function

The NumPy sort() function returns a sorted copy of the specified array.

Syntax

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

Parameters

Example:

In the example below, sort() function is used to sort elements of a 2-D array. As the axis parameter is not provided, the sorting is done along the last axis (row-wise).

import numpy as np
Arr = np.array([[1,20,5],[21, 4, 3],[11, 5, 50]])
SortedArr = np.sort(Arr)

print("Original Array:")
print(Arr)
print("\nSorted Array:")
print(SortedArr)

The output of the above code will be:

Original Array:
[[ 1 20  5]
 [21  4  3]
 [11  5 50]]

Sorted Array:
[[ 1  5 20]
 [ 3  4 21]
 [ 5 11 50]]

Example:

To sort the array column-wise, axis parameter can be set to 0. When axis=None is used, the array is flattened before sorting. Consider the example below:

import numpy as np
Arr = np.array([[1,20,5],[21, 4, 3],[11, 5, 50]])
SortedArr1 = np.sort(Arr, axis=0)
SortedArr2 = np.sort(Arr, axis=None)

print("Original Array:")
print(Arr)
print("\nSorted Array1:")
print(SortedArr1)
print("\nSorted Array2:")
print(SortedArr2)

The output of the above code will be:

Original Array:
[[ 1 20  5]
 [21  4  3]
 [11  5 50]]

Sorted Array1:
[[ 1  4  3]
 [11  5  5]
 [21 20 50]]

Sorted Array2:
[ 1  3  4  5  5 11 20 21 50]

numpy.argsort() function

The NumPy argsort() function returns the indices that would sort an array.

Syntax

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

Parameters

Example:

In the example below, argsort() function is used to get the indices which is further used to yield the sorted array.

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

#displaying the array
print("The original array:")
print(Arr)

#getting the indices from argsort()
x = np.argsort(Arr)
print("\nIndices from argsort():")
print(x)

#yielding sorted array
print("\nSorted array:")
print(Arr[x])

The output of the above code will be:

The original array:
[50 40 10 60 30 20]

Indices from argsort():
[2 5 4 1 0 3]

Sorted array:
[10 20 30 40 50 60]

Example:

The order parameter can be used to specify which field need to sort first. Consider the example below:

import numpy as np
Arr = np.array([(20,60),
                (20,50),
                (20,55),
                (10,75),
                (10,25),
                (10,50)], 
                dtype=[('x', '<i4'), ('y', '<i4')])

#displaying the array
print("The original array:")
print(Arr)

#yielding sorted array
indices = np.argsort(Arr, order=('x','y') )
print("\nSorted array:")
print(Arr[indices])

The output of the above code will be:

The original array:
[(20, 60) (20, 50) (20, 55) (10, 75) (10, 25) (10, 50)]

Sorted array:
[(10, 25) (10, 50) (10, 75) (20, 50) (20, 55) (20, 60)]

numpy.lexsort() function

The NumPy lexsort() function performs an indirect stable sort using a sequence of keys.

When multiple sorting keys are provided, it can be interpreted as columns, lexsort() returns an array of integer indices that describes the sort order by multiple columns. The last key in the sequence is used for the primary sort order, the second-to-last key for the secondary sort order, and so on. The keys argument must be a sequence of objects that can be converted to arrays of the same shape. If a 2D array is provided for the keys argument, its rows are interpreted as the sorting keys and sorting is according to the last row, second last row etc.

Syntax

numpy.lexsort(keys, axis=-1)

Parameters

Example:

In the example below, lexsort() function is used to sort by x column first then by y column.

import numpy as np

#x column - First column
x = np.array([10, 20, 10, 20, 10, 25, 10])
#y column - Second column
y = np.array([40, 10, 45, 60, 50, 25, 30])

#getting the array of indices that sorts
#x column first, y column second
indices = np.lexsort((y, x))

#displaying the indices
print("Array of indices to sort columns")
print(indices)

#using indices to sort columns
print("\nSorted x and y columns:")
for i in indices:
  print(x[i], y[i])

The output of the above code will be:

Array of indices to sort columns
[6 0 2 4 1 3 5]

Sorted x and y columns:
10 30
10 40
10 45
10 50
20 10
20 60
25 25

numpy.argmax() and numpy.argmin() functions

The NumPy argmax() and argmin() functions returns the indices of the maximum and minimum values along an axis respectively. It is calculated over the flattened array by default, otherwise over the specified axis.

Syntax

numpy.argmax(a, axis=None, out=None)
numpy.argmin(a, axis=None, out=None)

Parameters

Example:

In the example below, functions is used to find out index of the maximum/minimum value in the whole array.

import numpy as np
Arr = np.arange(12).reshape(3,4)

print("Array is:")
print(Arr)

#index of maximum value
idx1 = np.argmax(Arr)
print("\nIndex of maximum value is:", idx1)

#index of minimum value
idx2 = np.argmin(Arr)
print("Index of minimum value is:", idx2)

The output of the above code will be:

Array is:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

Index of maximum value is: 11
Index of minimum value is: 0

Example:

When axis parameter is provided, index of maximum/minimum value is calculated over the specified axes. Consider the following example.

import numpy as np
Arr = np.arange(12).reshape(3,4)

print("Array is:")
print(Arr)

#Index of maximum value along axis=0
print("\nIndex of maximum value along axis=0")
print(np.argmax(Arr, axis=0))

#Index of maximum value along axis=1
print("\nIndex of maximum value along axis=1")
print(np.argmax(Arr, axis=1))

#Index of minimum value along axis=0
print("\nIndex of minimum value along axis=0")
print(np.argmin(Arr, axis=0))

#Index of minimum value along axis=1
print("\nIndex of minimum value along axis=1")
print(np.argmin(Arr, axis=1))

The output of the above code will be:

Array is:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

Index of maximum value along axis=0
[2 2 2 2]

Index of maximum value along axis=1
[3 3 3]

Index of minimum value along axis=0
[0 0 0 0]

Index of minimum value along axis=1
[0 0 0]

numpy.where() function

The NumPy where() function returns elements chosen from x or y depending on condition. When only condition is provided, the function returns the indices of elements of the given array which satisfies the condition.

Syntax

numpy.where(condition, x, y)

Parameters

Example:

In the example below, where() function is used to replace all negative elements with 0 from an array.

import numpy as np

x = np.arange(-2, 5)

#replacing all negative elements with 0
y = np.where(x > 0, x, 0)

#displaying the content of x and y
print("x contains:", x)
print("y contains:", y)

The output of the above code will be:

x contains: [-2 -1  0  1  2  3  4]
y contains: [0 0 0 1 2 3 4]

Example:

In this example, where() function is used to choose elements from two array based on a given condition.

import numpy as np

x = np.asarray([[10, 20], [30, 40]])
y = np.asarray([[15, 15], [25, 25]])

#applying where condition
z = np.where(x > y, x, y)

#displaying the content of x, y and z
print("x =")
print(x)
print("\ny =")
print(y)
print("\nz =")
print(z)

The output of the above code will be:

x =
[[10 20]
 [30 40]]

y =
[[15 15]
 [25 25]]

z =
[[15 20]
 [30 40]]

numpy.nonzero() function

The NumPy nonzero() function returns the indices of the elements that are non-zero. Returned value is a tuple of arrays, one for each dimension of a, containing the indices of the non-zero elements in that dimension. The values in a are always tested and returned in row-major, C-style order.

Syntax

numpy.nonzero(a)

Parameters

Example:

In the example below, nonzero() function is used to find out index of all elements which are non-zero.

import numpy as np
Arr = np.array([[1, 2, 0], 
                [0, 5, 0], 
                [7, 0, 9]])

print("Array is:")
print(Arr)

#index of all elements which are non-zero
x = np.nonzero(Arr)
print("\nIndices of nonzero elements:")
print(x)

#displaying all nonzero elements
print("\nAll nonzero elements:")
print(Arr[x])

The output of the above code will be:

Array is:
[[1 2 0]
 [0 5 0]
 [7 0 9]]

Indices of nonzero elements:
(array([0, 0, 1, 2, 2]), array([0, 1, 1, 0, 2]))

All nonzero elements:
[1 2 5 7 9]

Example:

A condition can be applied to nonzero() function as shown in the example below.

import numpy as np
Arr = np.array([90, 80, 10, 20, 50])

print("Array is:")
print(Arr)

#index of all elements which are greater than 25
x = np.nonzero(Arr > 25)
print("\nIndices of elements greater than 25:")
print(x)

#displaying all elements greater than 25
print("\nAll elements greater than 25:")
print(Arr[x])

The output of the above code will be:

Array is:
[90 80 10 20 50]

Indices of elements greater than 25:
(array([0, 1, 4]),)

All elements greater than 25:
[90 80 50]

numpy.extract() function

The NumPy extract() function returns the elements of an array that satisfy some condition. If condition is boolean the function is equivalent to arr[condition].

Syntax

numpy.extract(condition, arr)

Parameters

Example:

In the example below, extract() function is used to extract elements from an array based on given condition.

import numpy as np

Arr = np.arange(12).reshape((3, 4))

#displaying the array
print("The original array:")
print(Arr)

#defining condition
condition = np.mod(Arr, 3) == 0
print("\nCondition is:")
print(condition)

#applying condition on array
print("\nExtracting elements based on condition:")
print(np.extract(condition, Arr))

The output of the above code will be:

The original array:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

Condition is:
[[ True False False  True]
 [False False  True False]
 [False  True False False]]

Extracting elements based on condition:
[0 3 6 9]