How to Use numpy.where Function [With Examples]

Numpy.where is a powerful tool for conditional element selection or transformation on NumPy arrays. It is commonly used in data processing, analysis, and algorithms. In this tutorial, we look at how you can use the numpy.where function and bring a few practical examples.

#What is numpy.where function?

Numpy, short for Numerical Python is a fundamental Python library used in data science. Numpy.where supports multi-dimensional arrays and is used to perform a wide range of mathematical computations on arrays, matrices, and vectors. The numpy.where() function returns elements of a numpy array where a specified condition is satisfied.

#numpy.where function syntax

The numpy.where() function takes the following format.

numpy.where(condition , outcome A, outcome B)

From the syntax above, you can see that the function takes three arguments.

The first argument is the condition to be evaluated, which is a boolean value.

outcome A, is returned when the condition evaluates to true.

outcome B, is returned when the condition evaluates to false.

Let us check out a few use cases of the numpy.where() function and see how it performs various operations on elements in an array.

#numpy.where with a single condition

The most basic numpy.where() function use case is filtering elements of an array with a single test condition. Consider the code below where the numpy.where() function returns elements of the array greater than or equal to 15.

import numpy as np

my_array = np.arange(21)

# returns elements of an array where the test condition is `true`, i.e elements equal to or greater than 15.

result = np.where(my_array  >= 15)

print(result)

Output

(array([15, 16, 17, 18, 19, 20], dtype=int64),)

The dtype=int64 simply tells us the data type of the array, which is a 64-bit array.

#numpy.where multiple conditions

The numpy.where() function can also take multiple conditions. You can use the following operators to carry out numpy.where multiple conditions:

The | operator for OR operations

The & operator for AND operations

In the following numpy.where example, when the condition (x > 2) & (x <= 8) evaluates to true, the new array includes the original values in the x array. Otherwise, the new array inserts a 0 value when the condition evaluates to false.

import numpy as np

x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

result = np.where((x > 2) & (x <= 8), x, 0)

print(result)

Output

[0 0 3 4 5 6 7 8 0 0]

In the next code example, when the condition (x < 4) | (x > 7) evaluates to true, the new array includes elements in the original array. Otherwise, the value 0 is inserted.

import numpy as np

x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

result = np.where((x <  4) & (x > 7), x, 0)

print(result)

Output

[ 1  2  3  0  0  0  0  8  9 10]

#numpy.where with two arrays

Suppose you have two one-dimensional arrays: x and y, as shown in the code snippet.

If the condition x > 2 holds true, the new array selects elements from the x array. Otherwise, if the condition is false, the new array selects elements from the y array.

import numpy as np

x = np.array([1, 2, 3, 4, 5, 6])
y = np.array([10, 20, 30, 40, 50, 60])

result = np.where(x > 2 , x , y)

print(result)

The new array will contain elements from the x array where the condition x > 2 is true and elements from the y array where the condition is false.

Output

[10 20  3  4  5  6]

#NOTE:

When passing all the 3 arguments to a numpy.where() function, all the arrays MUST be of the same length. Otherwise, you will run into a Value Error error.

#numpy.where with arithmetic operations

You can also perform arithmetic operations on array elements using the numpy.where() function. Consider the sample code shown below.

When the condition x > 0 evaluates to true, the new array includes the original elements from the original array. Otherwise, when the condition evaluates to false, the new array will return the square of the elements inside the original array.

import numpy as np

x = np.array([-5, 2, 3, -1, -3, 4, 6])

result = np.where(x > 0, x , x * x)

print(result)

Output

[25  2  3  1  9  4  6]

#Replace elements with numpy.where()

You can replace some of the elements in an array with desired ones. In the following code, you can observe that we have replaced all negative elements with 0.

import numpy as np

x = np.array([-5, 2, -3, -1, 3, 4, 6, 9, -1])
result = np.where(x > 0, x , 0)

print(result)

Output

[0 2 0 0 3 4 6 9 0]

#Conclusion

In this numpy.where function tutorial, we have looked at the numpy.where() function, its syntax, and parameters. We have seen how you can use it to test conditions on an array and generate a new array based on the test condition.