La función where

La función numpy.where acepta como argumentos de entrada una condición y dos estructuras tipo array, x e y, y devuelve valores de x o de y en función de que se cumpla o no la condición: Si se cumple, se devuelve el valor de x. Si no se cumple, el de y.

En este ejemplo partimos de dos sencillas listas (x e y) y otra (c) con la condición:

Función numpy.where. Ejemplo de uso

El resultado ha incluido el primer y tercer elemento de la lista x y el segundo de la lista y, tal y como indican los elementos de la lista c.

En el ejemplo visto hemos declarado explícitamente los valores de la condición pero, frecuentemente, ésta depende del valor de otras variables, o incluso de los valores de x o de y. Por ejemplo, supongamos que deseamos extraer los elementos de x cuando éstos sean mayores que 2, y los elementos de y en el resto de casos. Podríamos hacerlo con el siguiente código:

Función numpy.where. Ejemplo de uso

Obsérvese cómo, en este caso, las variables x e y han tenido que ser convertidas en arrays NumPy pues, de otra forma, la comparación x > 2 daría un error (la comparación no está soportada entre listas y escalares ¡esa es exactamente la funcionalidad ofrecida por NumPy!).

En la práctica, frecuentemente se obvia la creación de la variable intermedia c, definiéndose la condición directamente en la función where:

Función numpy.where. Ejemplo de uso

Podríamos interpretar esta última instrucción como "donde x sea mayor que 2, devuelve el valor de x. Si no, devuelve el valor de y".