Como Usar o Método numpy.where() do Python

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

Por exemplo, se a 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á os outros

Resultado

array([0 1 2])

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

Mas e se quisermos preservar a dimensão do resultado e não perder elementos de nossa matriz 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, o que isso diz é que se condition for verdadeira para algum elemento em nossa matriz, a nova matriz escolherá elementos de x.

Senão, se for falso, serão utilizados elementos de y.

Com isso, nossa matriz de saída final será uma matriz 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 especificar também y. Isso ocorre porque, neste caso, a forma da matriz de saída deve ser a mesma da matriz de entrada.

NOTA: A mesma lógica se aplica tanto para matrizes 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 esta função.

Usando Python numpy.where()

Suponhamos que desejamos selecionar 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 de 2 dimensões 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 você 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ê agora tentar executar o código acima, com essa mudança, você obterá uma saída como esta:

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

Se observarmos 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, esta função é na verdade equivalente a np.asarray.nonzero().

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

Portanto, vamos agora ver um exemplo mais simples, que nos mostra como podemos ser flexíveis 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 por 10. Isso é realmente o que obtemos!

Transmissão com numpy.where()

Se fornecermos todos os arrays 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 forma (3, 4), (3, 4) e (1, 4)
c = np.where(a < 5, a, b * 10)

print(c)

Saída

[[ 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, a saída é selecionada com base na condição, então todos os elementos, mas aqui, b é transmitido para a forma de a. (Uma de suas dimensões tem apenas um elemento, então não haverá erros durante a transmissão)

Então, b agora se tornará [[0 1 2 3] [0 1 2 3] [0 1 2 3]], e agora, podemos selecionar elementos mesmo deste array transmitido.

Portanto, a forma da saída é a mesma que a forma 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() em Python