NumPy - Sort, Search & Counting Functions

---

---

NumPy provides a variety of functions for sorting, searching, and counting elements in arrays. These functions can be extremely useful for data manipulation and analysis.

NumPy offers several sorting algorithms, each with its characteristics. Following is a comparison of three common sorting algorithms −

The numpy.sort() Function

The sort() function returns a sorted copy of the input array. It can sort arrays along any specified axis and supports different sorting algorithms. Following is the syntax −

numpy.sort(a, axis, kind, order)

Where,

Example

In the following example, we are sorting a 2D NumPy array both by default and along a specific axis. We also demonstrate sorting a structured array by a specific field, such as 'name' −

import numpy as np  

# Create a 2D array
a = np.array([[3, 7], [9, 1]])

print("Our array is:",a)

# Default sort
print("Applying sort() function:",np.sort(a))

# Sort along axis 0
print("Sort along axis 0:",np.sort(a, axis=0))

# Order parameter in sort function
dt = np.dtype([('name', 'S10'), ('age', int)])
a = np.array([("raju", 21), ("anil", 25), ("ravi", 17), ("amar", 27)], dtype=dt)

print("Our array is:",a)

print("Order by name:",np.sort(a, order='name'))

It will produce the following output −

Our array is:
[[3 7]
 [9 1]]

Applying sort() function:
[[3 7]
 [1 9]]

Sort along axis 0:
[[3 1]
 [9 7]]

Our array is:
[('raju', 21) ('anil', 25) ('ravi', 17) ('amar', 27)]

Order by name:
[('amar', 27) ('anil', 25) ('raju', 21) ('ravi', 17)]

The numpy.argsort() Function

The numpy.argsort() function performs an indirect sort on input array, along the given axis and using a specified kind of sort to return the array of indices of data. This indices array is used to construct the sorted array.

Example

In this example, we are retrieving indices, which are the positions of the sorted elements in the original array using the argsort() function. Using these indices, you can reconstruct the sorted array −

import numpy as np 

# Create an array
x = np.array([3, 1, 2])

print("Our array is:",x)

# Get indices that would sort the array
y = np.argsort(x)

print("Applying argsort() to x:",y)

# Reconstruct the sorted array using the indices
print("Reconstruct original array in sorted order:",x[y])

# Reconstruct the original array using a loop
print("Reconstruct the original array using loop:")
for i in y:
   print(x[i], end=' ')

It will produce the following output −

Our array is:
[3 1 2]

Applying argsort() to x:
[1 2 0]

Reconstruct original array in sorted order:
[1 2 3]

Reconstruct the original array using loop:
1 2 3

The numpy.lexsort() Function

The NumPy lexort() function performs an indirect sort using a sequence of keys. The keys can be seen as a column in a spreadsheet. The function returns an array of indices, using which the sorted data can be obtained. Note, that the last key happens to be the primary key of sort.

Example

In this example, we are using np.lexsort() function to sort a dataset based on multiple keys, where the last key "nm" is the primary sorting criterion. The sorted indices are then used to display the sorted data by combining the names and corresponding fields −

import numpy as np 

# Define keys
nm = ('raju', 'anil', 'ravi', 'amar') 
dv = ('f.y.', 's.y.', 's.y.', 'f.y.') 

# Get indices for sorted order
ind = np.lexsort((dv, nm))

print("Applying lexsort() function:",ind)

# Use indices to get sorted data
print("Use this index to get sorted data:",[nm[i] + ", " + dv[i] for i in ind]) 

It will produce the following output −

Applying lexsort() function:
[3 1 0 2]

Use this index to get sorted data:
['amar, f.y.', 'anil, s.y.', 'raju, f.y.', 'ravi, s.y.']

NumPy provides functions to find indices of maximum, minimum, and non-zero elements, as well as elements satisfying a condition.

The numpy.argmax() and numpy.argmin() Functions

The NumPy argmax() and argmin() functions return the indices of maximum and minimum elements respectively along the given axis.

Example

In this example, we are using np.argmax() and np.argmin() functions to find the indices of maximum and minimum values in a 2D array, both in the flattened array and along specific axes −

import numpy as np 

# Create a 2D array
a = np.array([[30, 40, 70], [80, 20, 10], [50, 90, 60]])

print("Our array is:",a)

# Apply argmax() function
print("Applying argmax() function:",np.argmax(a))

# Index of maximum number in flattened array
print("Index of maximum number in flattened array:",a.flatten())

# Array containing indices of maximum along axis 0
print("Array containing indices of maximum along axis 0:")
maxindex = np.argmax(a, axis=0)
print(maxindex)


# Array containing indices of maximum along axis 1
print("Array containing indices of maximum along axis 1:")
maxindex = np.argmax(a, axis=1)
print(maxindex)

# Apply argmin() function
print("Applying argmin() function:")
minindex = np.argmin(a)
print(minindex)

# Flattened array
print("Flattened array:",a.flatten()[minindex])

# Flattened array along axis 0
print("Flattened array along axis 0:")
minindex = np.argmin(a, axis=0)
print(minindex)

# Flattened array along axis 1
print("Flattened array along axis 1:")
minindex = np.argmin(a, axis=1)
print(minindex)

The output includes indices for these extrema, demonstrating how to access and interpret these positions within the array −

Our array is:
[[30 40 70]
 [80 20 10]
 [50 90 60]]

Applying argmax() function:
7

Index of maximum number in flattened array
[30 40 70 80 20 10 50 90 60]

Array containing indices of maximum along axis 0:
[1 2 0]

Array containing indices of maximum along axis 1:
[2 0 1]

Applying argmin() function:
5

Flattened array:
10

Flattened array along axis 0:
[0 1 1]

Flattened array along axis 1:
[0 2 0]

The numpy.nonzero() Function

The numpy.nonzero() function returns the indices of non-zero elements in the input array.

Example

In the example below, we are retrieving the indices of non-zero elements in the array "a" using the nonzero() function −

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

print ('Our array is:',a)
print ('Applying nonzero() function:',np.nonzero (a))

It will produce the following output −

Our array is:
[[30 40 0]
 [ 0 20 10]
 [50 0 60]]

Applying nonzero() function:
(array([0, 0, 1, 1, 2, 2]), array([0, 1, 1, 2, 0, 2]))

The numpy.where() Function

The where() function returns the indices of elements in an input array where the given condition is satisfied as shown in the example below −

import numpy as np 
x = np.arange(9.).reshape(3, 3) 

print ('Our array is:',x)  

print ('Indices of elements > 3')
y = np.where(x > 3) 
print (y)  

print ('Use these indices to get elements satisfying the condition',x[y])

It will produce the following output −

Our array is:
[[ 0. 1. 2.]
 [ 3. 4. 5.]
 [ 6. 7. 8.]]

Indices of elements > 3
(array([1, 1, 2, 2, 2]), array([1, 2, 0, 1, 2]))

Use these indices to get elements satisfying the condition
[ 4. 5. 6. 7. 8.]

The numpy.extract() Function

The extract() function returns the elements satisfying any condition as shown in the example below −

import numpy as np 
x = np.arange(9.).reshape(3, 3) 

print ('Our array is:',x)  

# define a condition 
condition = np.mod(x,2) == 0 

print ('Element-wise value of condition',condition)

print ('Extract elements using condition',np.extract(condition, x))

It will produce the following output −

Our array is:
[[ 0. 1. 2.]
 [ 3. 4. 5.]
 [ 6. 7. 8.]]

Element-wise value of condition
[[ True False True]
 [False True False]
 [ True False True]]

Extract elements using condition
[ 0. 2. 4. 6. 8.]