Numpy básico: seleccionar elementos condicionalmente en Numpy

Una tarea que se realiza casi a diario con Numpy es seleccionar elementos de un vector. Lo que se puede hacer en base a índices o, más interesante y productivo todavía, en base a una o varias condiciones. Como que los elementos a seleccionar sean menores que un valor o se encuentre en un rango. Esto es lo que se mostrará a continuación: cómo seleccionar elementos condicionalmente en Numpy.

A modo de ejemplo en el resto de la entrada se emplea un vector con los números entre 0 y 9, lo que se puede crear fácilmente con np.arrange.

import numpy as np

arr = np.arange(10)
arr

array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

Seleccionar elementos en base a una condición en Numpy

Al aplicar un operador de igualdad o comparación a un objeto Numpy se obtiene un nuevo vector de valores booleanos. Por ejemplo, si se utiliza el operador menor que para comparar el vector de ejemplo con un número como 5 se obtiene lo siguiente:

arr < 5

array([ True,  True,  True,  True,  True, False, False, False, False, False])

Un nuevo vector con valores booleanos en el que los primeros cinco elementos son ciertos, mientras que los últimos cinco son falsos. Este resultado se puede guardar en una variable y utilizar para seleccionar los elementos que cumple la condición. Si a un array de Numpy se le pasa un vector de valores booleanos el resultado es un nuevo vector solamente con elementos en las posiciones que son ciertas. Lo que se puede ver en el siguiente ejemplo.

menor_5 = arr < 5
arr[menor_5]

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

Aunque es más compacto no guardar la comparación en una variable, sino que utilizar directamente para seleccionar los elementos.

arr[arr < 5]

Seleccionar elementos en base a múltiples condiciones en Numpy

Una vez entendido que como se puede utilizar vectores booleanos para seleccionar elementos en Numpy se puede construir operaciones complejas que permitan seleccionar lo que se desee. Pudiendo así procesar los datos directamente en Python. Por ejemplo, nos podemos plantear como seleccionar los valores que son pares mayores de cinco. Los elementos pares serán aquellos cuyo resto de dividir por dos es cero y los mayores de cinco se obtienen con el operador mayor que. Esto es, el ejemplo se puede resolver con la siguiente línea.

arr[(np.mod(arr, 2) == 0) & (arr > 5)]

array([6, 8])

Conclusiones

En esta entrada se ha visto la forma de seleccionar elementos condicionalmente en Numpy. Una técnica con la que se mejorar la capacidad de analizar grandes conjuntos de datos en Python de una forma muy eficiente. Ya que facilita implementar fácilmente filtros complejos.

Imágenes: Pixabay (Marit Welker)