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 −
| kind | speed | worst case | work space | stable |
|---|---|---|---|---|
| ‘quicksort’ | 1 | O(n^2) | 0 | no |
| ‘mergesort’ | 2 | O(n*log(n)) | ~n/2 | yes |
| ‘heapsort’ | 3 | O(n*log(n)) | 0 | no |
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,
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’ −
# Open Compiler
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'))
Output:
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 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.
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 −
# Open Compiler
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=' ')
Output:
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 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.
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 −
# Open Compiler
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])
Output:
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 argmin() functions return the indices of maximum and minimum elements respectively along the given axis.
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 −
# Open Compiler
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)
Output:
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 returns the indices of non-zero elements in the input array.
In the example below, we are retrieving the indices of non-zero elements in the array “a” using the nonzero() function −
# Open Compiler
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))
Output:
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 where() function returns the indices of elements in an input array where the given condition is satisfied as shown in the example below −
# Open Compiler
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])
Output:
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 extract() function returns the elements satisfying any condition as shown in the example below −
# Open Compiler
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))
Output:
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.]
Key Takeaway: Master sorting, searching, and counting in NumPy with Vista Academy!