Cómo usar el método numpy.where() de Python

En Python, podemos usar la función numpy.where() para seleccionar elementos de un array de numpy, basándonos en una condición.

No solo eso, sino que también podemos realizar algunas operaciones en esos elementos si se cumple la condición.

¡Veamos cómo podemos usar esta función, utilizando algunos ejemplos ilustrativos!

Sintaxis de Python numpy.where()

Esta función acepta un array similar a numpy (por ejemplo, un array de NumPy de enteros/booleanos).

Devuelve un nuevo array de numpy, después de filtrar basado en una condición, que es un array similar a numpy de valores booleanos.

Por ejemplo, la condición puede tomar el valor de array([[True, True, True]]), que es un array booleano similar a numpy. (Por defecto, NumPy solo admite valores numéricos, pero también podemos convertirlos a bool)

Por ejemplo, si la condición es array([[True, True, False]]), y nuestro array es a = ndarray([[1, 2, 3]]), al aplicar una condición al array (a[:, condición]), obtendremos el array ndarray([[1 2]]).

import numpy as np

a = np.arange(10)
print(a[a <= 2]) # Solo capturará elementos <= 2 e ignorará los demás

Salida

array([0 1 2])

NOTA: La misma condición también se puede representar como a <= 2. Este es el formato recomendado para el arreglo de condiciones, ya que escribirlo como un arreglo booleano es muy tedioso

Pero ¿qué pasa si queremos preservar la dimensión del resultado y no perder elementos de nuestro arreglo original? Podemos usar numpy.where() para esto.

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

Tenemos dos parámetros más x y y. ¿Qué son esos?

Básicamente, lo que esto dice es que si condición es verdadera para algún elemento en nuestro arreglo, el nuevo arreglo elegirá elementos de x.

De lo contrario, si es falso, se tomarán elementos de y.

Con eso, nuestro arreglo de salida final será un arreglo con elementos de x donde condición = True, y elementos de y cuando condición = False.

Tenga en cuenta que aunque x y y son opcionales, si especifica x, DEBE especificar también y. Esto se debe a que, en este caso, la forma del arreglo de salida debe ser la misma que la del arreglo de entrada.

NOTA: La misma lógica se aplica tanto para arreglos unidimensionales como multidimensionales. En ambos casos, filtramos según la condición. Además, recuerde que las formas de x, y y condición se transmiten juntas.

Ahora, veamos algunos ejemplos para entender esta función correctamente.

Usando Python numpy.where()

Supongamos que queremos tomar solo los elementos positivos de un arreglo numpy y establecer todos los elementos negativos en 0, escribamos el código usando numpy.where().

1. Reemplazar elementos con numpy.where()

Usaremos aquí una matriz aleatoria de 2 dimensiones, y solo mostraremos los elementos positivos.

import numpy as np

# Inicialización aleatoria de un arreglo (2D)
a = np.random.randn(2, 3)
print(a)

# b será todos los elementos de a siempre que la condición sea verdadera (es decir, solo los elementos positivos)
# De lo contrario, establecerlo como 0
b = np.where(a > 0, a, 0)

print(b)

Salida Posible

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

¡Como puedes ver, ahora solo se retienen los elementos positivos!

2. Usando numpy.where() con solo una condición

Puede haber cierta confusión con el código anterior, ya que algunos de ustedes pueden pensar que la forma más intuitiva sería simplemente escribir la condición de esta manera:

import random
import numpy as np

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

Si ahora intentas ejecutar el código anterior, con este cambio, obtendrás una salida como esta:

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

Si observas detenidamente, b ahora es una tuple de arrays de numpy. Y cada array es la ubicación de un elemento positivo. ¿Qué significa esto?

Cuando proporcionamos solo una condición, esta función es en realidad equivalente a np.asarray.nonzero().

En nuestro ejemplo, np.asarray(a > 0) devolverá un array tipo booleano después de aplicar la condición, y np.nonzero(arr_like) devolverá los índices de los elementos no nulos de arr_like. (Consulta este enlace)

Entonces, ahora veremos un ejemplo más simple, que nos muestra cuán flexible podemos ser con numpy!

import numpy as np

a = np.arange(10)

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

print(a)
print(b)

Salida

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

Aquí, la condición es a < 5, que será el array tipo numpy [True True True True True False False False False False], x es el array a, y y es el array a * 10. Entonces, elegimos de a solo si a < 5, y de a * 10, si a > 5.

Entonces, esto transforma todos los elementos >= 5, multiplicándolos por 10. ¡Esto es lo que obtenemos realmente!

Transmisión con numpy.where()

Si proporcionamos todos los arrays condición, x y y, numpy los transmitirá juntos.

import numpy as np

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

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

print(a)
print(b)

# Transmisiones (a < 5, a y b * 10)
# de forma (3, 4), (3, 4) y (1, 4)
c = np.where(a < 5, a, b * 10)

print(c)

Salida

[[ 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]]

Nuevamente, aquí, la salida se selecciona según la condición, por lo que todos los elementos, pero aquí, b se transmite a la forma de a. (Una de sus dimensiones tiene solo un elemento, por lo que no habrá errores durante la transmisión)

Entonces, b ahora será [[0 1 2 3] [0 1 2 3] [0 1 2 3]], y ahora, podemos seleccionar elementos incluso de este array transmitido.

Por lo tanto, la forma de la salida es la misma que la forma de a.

Conclusión

En este artículo, aprendimos cómo podemos usar la función numpy.where() de Python para seleccionar arrays basados en otro array de condición.

Referencias

• Documentación de SciPy sobre la función numpy.where() en Python