Como usar np.where em Python

A função np.where() permite que você trabalhe es­pe­ci­fi­ca­mente com arrays NumPy na linguagem de pro­gra­ma­ção Python. Graças a suas operações ve­to­ri­za­das, 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 é pro­ve­ni­ente da bi­bli­o­teca NumPy, sendo adequada para a seleção con­di­ci­o­nal de elementos em um array. Ela iden­ti­fica e extrai elementos que atendem a uma condição es­pe­cí­fica, re­tor­nando os índices ou valores que cor­res­pon­dem a essa condição.

Em Python, np.where() pode ser usada em di­fe­ren­tes áreas, como no pro­ces­sa­mento de dados, em cálculos ci­en­tí­fi­cos, no apren­di­zado de máquina e em análises de dados. Na ma­ni­pu­la­ção de dados es­pe­ci­fi­ca­mente, np.where() pos­si­bi­lita fil­tra­gens de dados com base em critérios es­pe­cí­fi­cos, assim como a subs­ti­tui­ção de valores em arrays.

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

A função np.where() aceita um array se­me­lhante 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 de­ter­mi­nar que elementos devem ser se­le­ci­o­na­dos.
• x e y (opcionais): Se apenas condition for es­pe­ci­fi­cada, np.where() retorna os índices dos elementos que atendem à condição. Quando x e y são também es­pe­ci­fi­ca­dos, valores de x são re­tor­na­dos se a condição for ver­da­deira 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 (ver­da­deiro) ou False (false) no array de condições cor­res­ponde à seleção ou à filtragem do elemento cor­res­pon­dente no array original.

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

A função np.where() é ex­tre­ma­mente versátil, pois ela tem a ca­pa­ci­dade de realizar operações con­di­ci­o­nais em arrays, sendo altamente re­co­men­dada em ma­ni­pu­la­ções de dados.

Subs­ti­tuir elementos em array NumPy com np.where()

A função np.where() permite subs­ti­tuir elementos em um array NumPy com base em uma condição. Com ela, é possível trans­for­mar 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 iden­ti­fi­car elementos maiores que 3. Podemos fazer isso criando uma máscara con­di­ci­o­nal, 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 subs­ti­tuí­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 mo­di­fi­cado após a subs­ti­tui­ção de acordo com a condição. A saída mostra que os valores maiores que 3 foram subs­ti­tuí­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 subs­ti­tui­ção es­pe­ci­fi­ca­dos, 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 iden­ti­fica elementos em arr maiores que 2. Quando np.where(condition) é chamado sem valores de subs­ti­tui­ção, ele retorna uma tupla de índices com a condição arr > 2 sa­tis­feita. A saída é um array de índices no qual elementos de arr são maiores que 2.

Bro­ad­cas­ting com np.where()

O poderoso recurso Bro­ad­cas­ting do NumPy permite a re­a­li­za­ção de operações entre arrays com formas di­fe­ren­tes, desde que certas condições sejam atendidas. Quando arrays têm formas di­fe­ren­tes, o NumPy tenta expandi-los para torná-los com­pa­tí­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 exa­ta­mente, o Bro­ad­cas­ting permite a re­a­li­za­çã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 Bro­ad­cas­ting:

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

O array uni­di­men­si­o­nal row foi expandido para uma matriz de forma (3, 3) para cor­res­pon­der à forma de arr, repetindo-se ao longo das linhas. Em seguida, elemento por elemento de arr e row foram somados.