Sort, search and count

`sort` and `argsort`

The np.sort returns a sorted copy. The np.argsort returns, instead of the values, the positions that would sort the array, and this is the most versatile tool:

import numpy as np

v = np.array([30, 10, 50, 20, 40])
print(np.sort(v), v)
print(np.argsort(v))
print(v[np.argsort(v)[::-1]])

[10 20 30 40 50] [30 10 50 20 40]
[1 3 0 4 2]
[50 40 30 20 10]

With the positions you sort an array by the key of another: sorting names by grades, for example. For descending order, sort the negative:

nomes = np.array(["Ana", "Bia", "Caio"])
notas = np.array([7.5, 9.0, 6.0])
print(nomes[np.argsort(-notas)])

['Bia' 'Ana' 'Caio']

In two dimensions, the axis chooses the direction of the sorting:

m = np.array([[3, 1, 2], [9, 8, 7]])
print(np.sort(m, axis=1))
print(np.sort(m, axis=0))

[[1 2 3]
 [7 8 9]]
[[3 1 2]
 [9 8 7]]

Sort by multiple criteria

The np.lexsort sorts by multiple keys, and the last key of the tuple is the main one. To sort by last name and, in case of a tie, by age, age comes first:

sobrenome = np.array(["Silva", "Souza", "Silva", "Souza"])
idade = np.array([30, 25, 20, 40])
ordem = np.lexsort((idade, sobrenome))
print(ordem.tolist())

[2, 0, 1, 3]

Search in a sorted array

The np.searchsorted finds, by binary search (fast even with millions of elements), the position where each value would fit. It is the basis for grouping values into ranges, and the np.digitize does this directly:

ordenado = np.array([10, 20, 30, 40])
print(np.searchsorted(ordenado, [5, 20, 25, 99]))

idades = np.array([5, 17, 18, 40, 65, 90])
faixas = np.digitize(idades, [18, 65])
print(faixas)

[0 1 2 4]
[0 0 1 1 2 2]

The ranges of the digitize are: 0 under 18, 1 from 18 to 64 and 2 from 65 onwards.

Unique values and membership

a = np.array([3, 1, 2, 3, 1])
print(np.unique(a), np.isin(a, [1, 2]))
valores, primeira, contagem = np.unique(a, return_index=True, return_counts=True)
print(valores, primeira, contagem)

[1 2 3] [False  True  True False  True]
[1 2 3] [1 2 0] [2 1 2]

The np.isin asks, for each element, if it is in a set, and replaces the old np.in1d.

The k largest, without sorting everything

To find the 3 largest among millions of values, sorting everything is a waste. The np.argpartition puts the k desired elements at the front, without sorting them, in time proportional to the size of the array:

pontos = np.array([55, 91, 12, 78, 66, 99, 40])
topo3 = np.argpartition(-pontos, 3)[:3]
print(sorted(pontos[topo3].tolist(), reverse=True))

[99, 91, 78]

Stable sorts

If two elements are equal, a stable sort maintains the original order between them. To ensure this in the argsort, pass kind="stable". This matters when you sort in stages (first by one key, then by another).