Ordenar, buscar e contar

`sort` e `argsort`

O np.sort devolve uma cópia ordenada. O np.argsort devolve, em vez dos valores, as posições que ordenariam o array, e essa é a ferramenta mais versátil:

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]

Com as posições você ordena um array pela chave de outro: classificar nomes pelas notas, por exemplo. Para a ordem decrescente, ordene o negativo:

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

['Bia' 'Ana' 'Caio']

Em duas dimensões, o axis escolhe a direção da ordenação:

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]]

Ordenar por vários critérios

O np.lexsort ordena por várias chaves, e a última chave da tupla é a principal. Para ordenar por sobrenome e, em caso de empate, por idade, a idade vem primeiro:

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]

Buscar em um array ordenado

O np.searchsorted encontra, por busca binária (rápida mesmo em milhões de elementos), a posição onde cada valor se encaixaria. Ele é a base para agrupar valores em faixas, e o np.digitize faz isso diretamente:

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]

As faixas do digitize são: 0 abaixo de 18, 1 de 18 até 64 e 2 de 65 em diante.

Valores únicos e pertencimento

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]

O np.isin pergunta, para cada elemento, se ele está em um conjunto, e substitui o antigo np.in1d.

Os k maiores, sem ordenar tudo

Para achar os 3 maiores entre milhões de valores, ordenar tudo é desperdício. O np.argpartition coloca os k elementos desejados na frente, sem ordená-los, em tempo proporcional ao tamanho do 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]

Ordenações estáveis

Se dois elementos são iguais, uma ordenação estável mantém a ordem original entre eles. Para garantir isso no argsort, passe kind="stable". Isso importa quando você ordena em etapas (primeiro por uma chave, depois por outra).