×

NumPy Sort and Search functions

Sorting functions

The NumPy package provides various functions to perform sorting operations on the numpy arrays.

  • sort()
  • argsort()
  • lexsort()
  • sort_complex()

sort():

The sort() function of NumPy sorts the input array. The sort() function returns a sorted copy of the input array with the same dimensions and shape as the input array.

Syntax:

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

The a parameter is the input array to be sorted, the axis parameter is the axis along which to sort (if not mentioned, the array is flattened before sorting), and the kind parameter is the sorting algorithm(default is quicksort). The axis default is -1, which sorts along the last axis.

import numpy as np
a = np.array([[4, 7, 2], [3, 8, 20], [23, 1, 12]])

print(np.sort(a, axis = 0, kind = 'heapsort'))
#Output:
[[ 3  1  2]
 [ 4  7 12]
 [23  8 20]]

print(np.sort(a, axis = 1, kind = 'mergesort'))
#Output:
[[ 2  4  7]
 [ 3  8 20]
 [ 1 12 23]]

argsort():

The argsort() function returns the indices of the sorted array. The argsort() function does an indirect sort along the given axis using the algorithm specified and returns an array of indices of the same shape as the input array.

Syntax:

numpy.argsort(a, axis = -1, kind=None)
x = np.array([0, 3, 4, 2])
indices = np.argsort(x, kind ='stable')

print(indices) #Sorted indices
#Output:
[0 3 1 2]

print(x[indices]) #sorted array
#Output:
[0 2 3 4]

lexsort():

The lexsort() function performs an indirect stable sort using a sequence of keys. The lexsort() function returns an array of integer indices that describes the sort order by multiple columns.

Syntax:

lexsort(keys, axis = -1)

The keys parameter is the array or tuple containing different “columns” to be sorted. The last column (or row if keys are a 2D array) is the primary sort key. The axis parameter is the axis to be sorted(By default, sort over the last axis).

a = np.array([2,5,2,4,3,4,4]) 
b = np.array([9,4,0,4,1,3,2])
indices = np.lexsort((b,a)) #sorting by a, then by b
res = []
for i in indices:
    res.append((a[i], b[i]))
print(indices) #sorted indices values for both arrays

#Output:
[2 0 4 6 5 3 1]

print(res) #sorted array

#Output:
[(2, 0), (2, 9), (3, 1), (4, 2), (4, 3), (4, 4), (5, 4)]

sort_complex():

The sort_complex() function sorts a complex array using the real part first, then the imaginary part.

Syntax:

sort_complex(a)

The a parameter is the input array.

c = np.array([[1 + 7j, 6 - 2j, 6 - 11j], [8 - 9j, 2 + 5j, 10 + 3j]])
print(np.sort_complex(c))

#Output:
[[ 1. +7.j  6.-11.j  6. -2.j]
 [ 2. +5.j  8. -9.j 10. +3.j]]

Searching functions

The NumPy package has several functions to search and get a specific portion of the numpy arrays.

  • argmax()
  • argmin()
  • where()
  • extract()

argmax():

The argmax() returns the array of indices of the maximum values in the input array along an axis. In the case of multiple occurrences of the maximum values, the indices corresponding to the first occurrence are returned.

Syntax:

numpy.argmax(a, axis = None)

The a parameter is the input array, the axis parameter is the axis along which to sort (By default, the index is into the flattened array, otherwise along the specified axis).

a = np.array([[0, 7, 5], [9, 8, 20], [2, 1, 3]])
print(np.argmax(a, axis = 1))

#Output:
[1 2 2]

argmin():

The argmin() gets the array of indices of the minimum values in the input array along an axis.

Syntax:

numpy.argmin(a, axis = None)
a = np.array([[0, 7, 5], [9, 8, 20], [2, 1, 3]])
print(np.argmin(a, axis = 1))

#Output:
[0 1 1]

where():

The where() function returns elements chosen from the two input arrays depending on the condition.

Syntax:

where(condition[, x, y])

The condition parameter is the condition when True, returns the first array element, otherwise returns the second array element), and the x and y parameters are the two input arrays.

a = np.arange(2, 8)
print(np.where(a<4,2*a,10+a))

#Output:
[ 4  6 14 15 16 17]

extract():

The extract() function gets the elements of an array that satisfy some condition.

Syntax:

extract(condition, arr)

The condition parameter is the condition to be satisfied for an extracted subarray, and the arr parameter is the input array.

a = np.array([[10, 27, 35], [19, 48, 70], [92, 51, 83]])
condition = np.mod(a, 5) == 0
print(np.extract(condition, a))

#Output:
[10 35 70]