Sort, Search & Counting Functions

NumPy - Sorting, Searching, and Counting Functions

NumPy offers a suite of efficient functions for sorting arrays, searching for elements, and counting occurrences.

Sorting Arrays

NumPy provides several sorting algorithms (quicksort, mergesort, heapsort, stable). The choice can impact performance and stability (whether the relative order of equal elements is preserved).

kind	Average Speed	Worst Case	Work Space	Stable
’quicksort’ (default)	Fastest	O(n^2)	Constant	No
’mergesort’	Medium	O(n*log(n))	~n/2	Yes
’heapsort’	Slowest	O(n*log(n))	Constant	No
’stable’	Medium	O(n*log(n))	~n/2	Yes (adapts mergesort/timsort)

Note: ‘stable’ is often a good choice when preserving the original order of equal elements matters.

numpy.sort()

Returns a sorted copy of an array. It does not modify the original array.

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

Parameters:

Parameter	Description
a	Input array to be sorted.
axis	Axis along which to sort. Default is -1 (the last axis). If None, the array is flattened before sorting.
kind	Sorting algorithm (‘quicksort’, ‘mergesort’, ‘heapsort’, ‘stable’). Default depends on NumPy version, often ‘quicksort’.
order	For structured arrays, specifies the field(s) to sort by.

Example: Using np.sort

Sorting a 2D array along different axes and sorting a structured array.

import numpy as np

# --- Basic Sorting ---
a = np.array([[3, 7], [9, 1]])

print('Original array:')
print(a)

# Sort along the last axis (rows) - default behavior
print('\nSorted along axis=-1 (rows):')
print(np.sort(a))

# Sort along axis 0 (columns)
print('\nSorted along axis=0 (columns):')
print(np.sort(a, axis=0))

# Sort flattened array
print('\nSorted flattened array:')
print(np.sort(a, axis=None))


# --- Sorting Structured Array ---
# Define a structured data type
dt = np.dtype([('name', 'S10'), ('age', int)])
people = np.array([('Raju', 21), ('Anil', 25), ('Ravi', 17), ('Amar', 27)], dtype=dt)

print('\n\nOriginal structured array:')
print(people)

# Sort by the 'name' field
print('\nSorted by name:')
print(np.sort(people, order='name'))

# Sort by the 'age' field
print('\nSorted by age:')
print(np.sort(people, order='age'))

Output:

Original array:
[[3 7]
 [9 1]]

Sorted along axis=-1 (rows):
[[3 7]
 [1 9]]

Sorted along axis=0 (columns):
[[3 1]
 [9 7]]

Sorted flattened array:
[1 3 7 9]


Original structured array:
[(b'Raju', 21) (b'Anil', 25) (b'Ravi', 17) (b'Amar', 27)]

Sorted by name:
[(b'Amar', 27) (b'Anil', 25) (b'Raju', 21) (b'Ravi', 17)]

Sorted by age:
[(b'Ravi', 17) (b'Raju', 21) (b'Anil', 25) (b'Amar', 27)]

Note: To sort an array in-place (modifying the original), use the ndarray.sort() method: a.sort().

numpy.argsort()

Performs an indirect sort. Instead of returning the sorted array, it returns the indices that would sort the original array. This is useful when you need to sort multiple related arrays based on the order of one.

Example: Using argsort

Finding indices to sort an array and using them for reconstruction.

import numpy as np

x = np.array([30, 10, 20])

print('Original array:')
print(x)

print('\nIndices that would sort x:')
indices = np.argsort(x)
print(indices)

print('\nReconstruct sorted array using indices:')
print(x[indices]) # Use fancy indexing with the indices

Output:

Original array:
[30 10 20]

Indices that would sort x:
[1 2 0]

Reconstruct sorted array using indices:
[10 20 30]

numpy.lexsort()

Performs an indirect stable sort using a sequence of keys (arrays). It sorts based on the last key provided first, then uses the second-to-last key to break ties, and so on. Returns the indices that sort the data based on these keys.

Example: Using lexsort

Sorting names based first on division (‘dv’) then by name (‘nm’).

import numpy as np

# Data arrays (keys for sorting)
names = np.array(['Raju', 'Anil', 'Ravi', 'Amar'])
divisions = np.array(['f.y.', 's.y.', 's.y.', 'f.y.'])

# Sort first by names, then by divisions (keys are provided last-to-first)
# So, primary sort key is 'names', secondary is 'divisions'
indices = np.lexsort((divisions, names))

print('Indices from lexsort (sort by names, then divisions):')
print(indices)

print('\nUse indices to get sorted data:')
print([f'{names[i]}, {divisions[i]}' for i in indices])

Output:

Indices from lexsort (sort by names, then divisions):
[3 1 0 2]

Use indices to get sorted data:
['Amar, f.y.', 'Anil, s.y.', 'Raju, f.y.', 'Ravi, s.y.']

Searching Arrays

NumPy provides functions to find maximum/minimum values or elements satisfying specific conditions.

numpy.argmax() and numpy.argmin()

Return the indices of the maximum and minimum elements along a given axis, respectively. If the array is flattened (axis=None), the index into the flattened array is returned.

Example: Using argmax and argmin

import numpy as np

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

print('Original array:')
print(a)

print('\nIndex of max element (flattened):', np.argmax(a))
print('Value of max element:', a.flat[np.argmax(a)])

print('\nIndices of max elements along axis 0 (columns):', np.argmax(a, axis=0))
print('Indices of max elements along axis 1 (rows):', np.argmax(a, axis=1))

print('\nIndex of min element (flattened):', np.argmin(a))
print('Value of min element:', a.flat[np.argmin(a)])

print('\nIndices of min elements along axis 0 (columns):', np.argmin(a, axis=0))
print('Indices of min elements along axis 1 (rows):', np.argmin(a, axis=1))

Output:

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

Index of max element (flattened): 7
Value of max element: 90

Indices of max elements along axis 0 (columns): [1 2 0]
Indices of max elements along axis 1 (rows): [2 0 1]

Index of min element (flattened): 5
Value of min element: 10

Indices of min elements along axis 0 (columns): [0 1 1]
Indices of min elements along axis 1 (rows): [0 2 0]

numpy.nonzero()

Returns the indices of elements that are non-zero. The result is a tuple of arrays, one array per dimension, containing the indices of the non-zero elements in that dimension.

Example: Using nonzero

import numpy as np

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

print('Original array:')
print(a)

print('\nIndices of non-zero elements:')
nonzero_indices = np.nonzero(a)
print(nonzero_indices)

# Access the non-zero elements using these indices
print('\nNon-zero elements:')
print(a[nonzero_indices])

Output:

Original array:
[[30 40  0]
 [ 0 20 10]
 [50  0 60]]

Indices of non-zero elements:
(array([0, 0, 1, 1, 2, 2]), array([0, 1, 1, 2, 0, 2]))

Non-zero elements:
[30 40 20 10 50 60]

numpy.where()

Returns indices where a given condition is true. Similar output format to nonzero. Can also be used like a ternary operator: np.where(condition, x, y) returns elements from x where condition is true, and elements from y where it’s false.

Example: Using where

import numpy as np

x = np.arange(9.).reshape(3, 3)

print('Original array:')
print(x)

condition = x > 3
print('\nCondition (x > 3):')
print(condition)

print('\nIndices where condition is true:')
indices = np.where(condition)
print(indices)

print('\nElements where condition is true:')
print(x[indices])

# Using where like a ternary operator
result = np.where(condition, x, -1) # Replace elements <= 3 with -1
print('\nArray with elements <= 3 replaced by -1:')
print(result)

Output:

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

Condition (x > 3):
[[False False False]
 [False  True  True]
 [ True  True  True]]

Indices where condition is true:
(array([1, 1, 2, 2, 2]), array([1, 2, 0, 1, 2]))

Elements where condition is true:
[4. 5. 6. 7. 8.]

Array with elements <= 3 replaced by -1:
[[-1. -1. -1.]
 [-1.  4.  5.]
 [ 6.  7.  8.]]

numpy.extract()

Returns the elements of an array that satisfy a given condition. It flattens the array before applying the condition.

Example: Using extract

import numpy as np

x = np.arange(9.).reshape(3, 3)

print('Original array:')
print(x)

# Define a condition (e.g., elements are even integers)
condition = np.mod(x, 2) == 0

print('\nCondition (x % 2 == 0):')
print(condition)

print('\nExtract elements using condition:')
print(np.extract(condition, x))

Output:

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

Condition (x % 2 == 0):
[[ True False  True]
 [False  True False]
 [ True False  True]]

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