Numpy básico: el método numpy.where()

En esta entrada se va a explicar el funcionamiento del método numpy.where() con el que se puede seleccionar elementos en base a una condición. Simplemente en una línea de código. El método numpy.where() puede evitar la necesidad de crear un bucle para crear un nuevo vector en base a una condición. Permitiendo así crear código más compacto.

Uso básico de numpy.where()

La forma más fácil de utilizar el método numpy.where() es como se muestra a continuación.

import numpy as np

np.where([True, False, True, False], [1, 2, 3, 4], [5, 6, 7, 8])

array([1, 6, 3, 8])

El primer parámetro es un vector de valores lógicos, verdadero o falso. A continuación, se inyectan dos vectores. La salida de la función es un nuevo vector con los elementos del primer o segundo vector en base a lo que indique el primer parámetro. El método funciona de la siguiente manera: si la condición para el elemento i es cierta se selecciona el elemento correspondiente del primer vector, en caso contrario del segundo. Esto es lo que se puede ver en ejemplo, en el primer y tercer caso se selecciona el elemento del primer vector, mientras que el resto de los casos del segundo.

Un ejemplo trivial de uso del método numpy.where() puede ser la creación de un nuevo vector con los registros más grandes de otros dos. Esto se puede conseguir como se muestra a continuación.

data_1 = np.array([1, 3, 2, 1])
data_2 = np.array([3, 2, 1, 3])

np.where(data_1 > data_2, data_1, data_2)

array([3, 3, 2, 3])

Los elementos sobre los que se itera no tienen porqué ser necesariamente vectores, también pueden ser un escalar. Lo que se puede utilizar para seleccionar para poner un valor en ciertas condiciones. Por ejemplo, para reemplazar los valores por debajo de un mínimo por dicho valor.

np.where(data_1 > 1, data_1, 2)

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

Uso avanzado de numpy.where()

Los ejemplos vistos hasta ahora son triviales, pero la condición lógica puede ser todo lo compleja que se desee. Pudiéndose utilizar cualquier operador lógico. Por ejemplo, a continuación, se muestra como seleccionar los elementos del primer vector cuando estos se encuentra en un rango y en caso contrario los del otro.

data_1 = np.array([9, 17, 15, 20])
data_2 = np.array([20, 5, 13, 18])

np.where((data_1 > 10) & (data_1 < 15), data_1, data_2)

array([20,  5, 13, 18])

Pudiendo aumentar la complejidad de las condiciones tanto como sea necesario en cada uno de los casos.

Conclusiones

El método numpy.where() de Numpy es una excelente herramienta para reemplazar los bucles a la hora de aplicar condiciones a vectores. Permitiendo así tener crear un código más compacto y legible.

Imágenes: Pixabay (Marit Welker)