Como usar o método Python numpy.where()

Em Python, podemos usar a função numpy.where() para selecionar elementos de um array numpy, com base em uma condição.

Não apenas isso, mas também podemos realizar algumas operações nesses elementos se a condição for satisfeita.

Vamos ver como podemos usar essa função, usando alguns exemplos ilustrativos!

Sintaxe do Python numpy.where()

Esta função aceita um array semelhante ao numpy (ex. um array NumPy de inteiros/booleanos).

Retorna um novo array numpy, após filtragem com base em uma condição, que é um array semelhante ao numpy de valores booleanos.

Por exemplo, condição pode ter o valor de array([[True, True, True]]), que é um array booleano semelhante ao numpy. (Por padrão, o NumPy só suporta valores numéricos, mas também podemos convertê-los para bool)

Por exemplo, se condição for array([[True, True, False]]), e nosso array for a = ndarray([[1, 2, 3]]), ao aplicar uma condição ao array (a[:, condição]), obteremos o array ndarray([[1 2]]).

import numpy as np

a = np.arange(10)
print(a[a <= 2]) # Capturará apenas elementos <= 2 e ignorará outros

Saída

array([0 1 2])

NOTA: A mesma condição também pode ser representada como a <= 2. Este é o formato recomendado para o array de condição, pois é muito tedioso escrevê-lo como um array booleano

Mas e se quisermos preservar a dimensão do resultado e não perder elementos do nosso array original? Podemos usar numpy.where() para isso.

numpy.where(condition [, x, y])

Temos mais dois parâmetros x e y. O que são esses?

Basicamente, isso significa que se condition for verdadeira para algum elemento em nosso array, o novo array escolherá elementos de x.

Se for falso, elementos de y serão escolhidos.

Com isso, nosso array de saída final será um array com elementos de x sempre que condition = True, e elementos de y sempre que condition = False.

Observe que, embora x e y sejam opcionais, se você especificar x, você DEVE também especificar y. Isso ocorre porque, neste caso, a forma do array de saída deve ser a mesma do array de entrada.

NOTA: A mesma lógica se aplica tanto para arrays unidimensionais quanto multidimensionais. Em ambos os casos, filtramos com base na condição. Lembre-se também de que as formas de x, y e condition são transmitidas juntas.

Agora, vamos ver alguns exemplos para entender melhor essa função.

Usando Python numpy.where()

Suponha que desejamos pegar apenas os elementos positivos de um array numpy e definir todos os elementos negativos como 0, vamos escrever o código usando numpy.where().

1. Substituir Elementos com numpy.where()

Vamos usar um array aleatório 2D aqui, e apenas mostrar os elementos positivos.

import numpy as np

# Inicialização aleatória de um array 2D
a = np.random.randn(2, 3)
print(a)

# b será todos os elementos de a sempre que a condição for verdadeira (ou seja, apenas elementos positivos)
# Caso contrário, defina como 0
b = np.where(a > 0, a, 0)

print(b)

Resultado Possível

[[-1.06455975  0.94589166 -1.94987123]
 [-1.72083344 -0.69813711  1.05448464]]
[[0.         0.94589166 0.        ]
 [0.         0.         1.05448464]]

Como pode ver, apenas os elementos positivos são mantidos agora!

2. Usando numpy.where() apenas com uma condição

Pode haver alguma confusão em relação ao código acima, pois alguns de vocês podem pensar que a maneira mais intuitiva seria simplesmente escrever a condição assim:

import random
import numpy as np

a = np.random.randn(2, 3)
b = np.where(a > 0)
print(b)

Se você tentar executar o código acima agora, com essa mudança, você obterá uma saída como esta:

(array([0, 1]), array([2, 1]))

Se você observar atentamente, b agora é uma tupla de matrizes numpy. E cada matriz é a localização de um elemento positivo. O que isso significa?

Sempre que fornecemos apenas uma condição, essa função é na verdade equivalente a np.asarray.nonzero().

No nosso exemplo, np.asarray(a > 0) retornará uma matriz booleana após aplicar a condição, e np.nonzero(arr_like) retornará os índices dos elementos não nulos de arr_like. (Consulte este link)

Então, agora vamos olhar para um exemplo mais simples, que nos mostra o quão flexíveis podemos ser com o numpy!

import numpy as np

a = np.arange(10)

b = np.where(a < 5, a, a * 10)

print(a)
print(b)

Saída

[0 1 2 3 4 5 6 7 8 9]
[ 0  1  2  3  4 50 60 70 80 90]

Aqui, a condição é a < 5, que será a matriz numpy [True True True True True False False False False False], x é a matriz a, e y é a matriz a * 10. Então, escolhemos de a apenas se a < 5, e de a * 10, se a > 5.

Portanto, isso transforma todos os elementos >= 5, multiplicando-os por 10. Isto é o que obtemos realmente!

Transmitindo com numpy.where()

Se fornecermos todos os arrays de condição, x e y, o numpy irá transmiti-los juntos.

import numpy as np

a = np.arange(12).reshape(3, 4)

b = np.arange(4).reshape(1, 4)

print(a)
print(b)

# Transmissões (a < 5, a e b * 10)
# de formato (3, 4), (3, 4) e (1, 4)
c = np.where(a < 5, a, b * 10)

print(c)

Resultado

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
[[0 1 2 3]]
[[ 0  1  2  3]
 [ 4 10 20 30]
 [ 0 10 20 30]]

Novamente, aqui, o resultado é selecionado com base na condição, então todos os elementos, mas aqui, b é transmitido para o formato de a. (Uma de suas dimensões tem apenas um elemento, então não haverá erros durante a transmissão)

Portanto, b agora se tornará [[0 1 2 3] [0 1 2 3] [0 1 2 3]], e agora, podemos selecionar elementos até mesmo deste array transmitido.

Então o formato do resultado é o mesmo que o formato de a.

Conclusão

Neste artigo, aprendemos como podemos usar a função numpy.where() do Python para selecionar arrays com base em outro array de condição.

Referências

• Documentação do SciPy sobre a função numpy.where() do Python