Como usar np.where em Python

A função np.where() permite que você trabalhe especificamente com arrays NumPy na linguagem de programação Python. Graças a suas operações vetorizadas, np.where consegue ser mais eficiente que métodos baseados em loops.

Para que serve a função np.where em Python?

A poderosa função np.where() em Python é proveniente da biblioteca NumPy, sendo adequada para a seleção condicional de elementos em um array. Ela identifica e extrai elementos que atendem a uma condição específica, retornando os índices ou valores que correspondem a essa condição.

Em Python, np.where() pode ser usada em diferentes áreas, como no processamento de dados, em cálculos científicos, no aprendizado de máquina e em análises de dados. Na manipulação de dados especificamente, np.where() possibilita filtragens de dados com base em critérios específicos, assim como a substituição de valores em arrays.

Sintaxe da função Python np.where()

A função np.where() aceita um array semelhante ao NumPy, que pode consistir, por exemplo, em inteiros ou valores booleanos. A sintaxe de np.where() em Python é a seguinte:

import numpy as np
np.where(condition[, x, y])

• condition: Condição aplicada ao array para determinar que elementos devem ser selecionados.
• x e y (opcionais): Se apenas condition for especificada, np.where() retorna os índices dos elementos que atendem à condição. Quando x e y são também especificados, valores de x são retornados se a condição for verdadeira e valores de y se a condição for falsa.

Quando a função np.where() é executada, ela retorna um novo array NumPy. Esse novo array é criado por meio da filtragem ou da seleção de elementos do array original, com base em uma condição. Cada valor True (verdadeiro) ou False (false) no array de condições corresponde à seleção ou à filtragem do elemento correspondente no array original.

Exemplos de aplicação da função Python np.where()

A função np.where() é extremamente versátil, pois ela tem a capacidade de realizar operações condicionais em arrays, sendo altamente recomendada em manipulações de dados.

Substituir elementos em array NumPy com np.where()

A função np.where() permite substituir elementos em um array NumPy com base em uma condição. Com ela, é possível transformar listas em Python em um array:

import numpy as np
arr = np.array([1, 2, 3, 4, 5])

Agora, vamos definir uma condição para identificar elementos maiores que 3. Podemos fazer isso criando uma máscara condicional, que dará True para elementos maiores que 3 e False para elementos menores que 3:

condition = arr > 3

Em seguida, aplicamos a função np.where() e passamos essa condição. Definimos que os valores que atenderem à condição devem ser substituídos por -1, enquanto todos os outros valores devem ser trocados por 0:

new_arr = np.where(condition, -1, 0)

O resultado, new_arr, contém o array modificado após a substituição de acordo com a condição. A saída mostra que os valores maiores que 3 foram substituídos por -1, enquanto os outros valores foram trocados por 0.

print(new_arr)
# Output: [0 0 0 -1 -1]

np.where()com somente uma condição

Se np.where() for usado apenas com uma condição e sem valores de substituição especificados, ele retorna uma tupla de índices com a condição True.

import numpy as np
arr = np.array([1, 2, 3, 4, 5])
condition = arr > 2
result = np.where(condition)
print(result)
# Output: (array([2, 3, 4]),)

Nesse exemplo, arr é um array NumPy contendo valores de 1 a 5. Por sua vez, condition = arr > 2 cria uma máscara booleana que identifica elementos em arr maiores que 2. Quando np.where(condition) é chamado sem valores de substituição, ele retorna uma tupla de índices com a condição arr > 2 satisfeita. A saída é um array de índices no qual elementos de arr são maiores que 2.

Broadcasting com np.where()

O poderoso recurso Broadcasting do NumPy permite a realização de operações entre arrays com formas diferentes, desde que certas condições sejam atendidas. Quando arrays têm formas diferentes, o NumPy tenta expandi-los para torná-los compatíveis.

No exemplo abaixo, temos um array NumPy arr com a forma (3, 3):

import numpy as np
arr = np.array([[1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]])

E um outro array NumPy row com a forma (3,):

row = np.array([10, 11, 12])

Observe que o array arr tem a forma (3, 3) e o array row tem a forma (3,). Embora as formas não coincidam exatamente, o Broadcasting permite a realização de uma operação entre esses arrays. Para adicionar os valores de cada linha de arr aos valores de row, por exemplo, poderemos lançar mão do Broadcasting:

result = arr + row
print(result)
# Output: [[11 13 15]
    [14 16 18]
    [17 19 21]]

O array unidimensional row foi expandido para uma matriz de forma (3, 3) para corresponder à forma de arr, repetindo-se ao longo das linhas. Em seguida, elemento por elemento de arr e row foram somados.