Indexação booleana e fancy

Uma condição vira uma máscara

Comparar um array com um valor produz outro array, de booleanos, chamado de máscara. Usá-la como índice filtra os elementos verdadeiros:

import numpy as np

notas = np.array([78, 92, 45, 88, 91, 30, 85])
mascara = notas > 75
print(mascara)
print(notas[mascara])
print(notas[notas > 75])
print((notas > 75).sum(), (notas > 75).mean().round(3))

[ True  True False  True  True False  True]
[78 92 88 91 85]
[78 92 88 91 85]
5 0.714

Como True vale 1 e False vale 0, o sum() de uma máscara conta quantos elementos passam, e o mean() dá a proporção. Aqui, 5 de 7 notas (71,4%) passaram de 75. É um dos truques mais usados em análise de dados.

Várias condições: `&`, `|` e `~`

Para combinar condições, use os operadores & (e), | (ou) e ~ (não), e sempre coloque cada condição entre parênteses, porque esses operadores têm precedência maior do que a comparação:

presenca = np.array([95, 60, 80, 92, 70, 88, 97])
print(notas[(notas > 80) & (presenca > 90)])
print(notas[(notas < 50) | (presenca < 65)])
print(notas[~(notas > 80)])

[88 85]
[92 45 30]
[78 45 30]

Usar and e or do Python é o erro clássico, e o NumPy responde com uma mensagem que parece enigmática:

try:
    notas[(notas > 80) and (presenca > 90)]
except ValueError as erro:
    print(erro)

The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

O and do Python precisa decidir se um array inteiro é verdadeiro ou falso, e um array com vários elementos não tem resposta única. É por isso que existem o & e o |, que operam elemento a elemento.

Indexação fancy: uma lista de posições

Passando uma lista de índices, você escolhe as posições que quiser, na ordem que quiser, inclusive repetindo:

nums = np.array([10, 20, 30, 40, 50])
print(nums[[0, 2, 4]])
print(nums[[4, 4, 0]])

tabela = np.arange(12).reshape(3, 4)
print(tabela[[0, 2], :])
print(tabela[[0, 1], [1, 3]])

[10 30 50]
[50 50 10]
[[ 0  1  2  3]
 [ 8  9 10 11]]
[1 7]

Com duas listas, o NumPy as combina par a par: o último resultado é tabela[0, 1] e tabela[1, 3], e não um bloco 2×2.

Máscara e fancy devolvem cópias

Diferente da fatia, a indexação booleana e a fancy sempre devolvem uma cópia. Alterar o resultado não toca o original:

recorte = nums[[0, 1]]
recorte[0] = 999
print(nums[0], np.shares_memory(nums, recorte))

10 False