NumPy where()

import numpy as np

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

condition = originalArray < 0


# For each element of originalArray,
# if condition is True, use 0 as element in the resultant array
# if condition is False, use the corresponding element in the resultant array
result = np.where(condition, 0, originalArray )

print(result)

# Output :  [1 0 0 4 5]

where() Syntax

The syntax of where() is:

numpy.where(condition, x, y)

where() Arguments

The where() method takes three arguments:

• condition - a boolean or an array
• x - value to take if the condition is True
• y - value to take if the condition is False

Note: We can also pass a single argument to np.where(). To learn about it, visit np.where() with a single argument section below.

where() Return Value

The where() method returns a new NumPy array.

Example 1: numpy.where() With Two Arrays

import numpy as np

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

test_condition = x < 3 

# if test_condition is True, select element of x
# if test_condition is False, select element of y
result = np.where(test_condition, x, y)

print(result)

Output

[1 2 30 40]

Example 2: numpy.where() with Operation

We can also use numpy.where() to perform operations on array elements.

import numpy as np

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

# test condition
test_condition = x > 0

# if test_condition is True, select element of x
# if test_condition is False, select  x * -1
result = np.where(test_condition, x, x * -1)

print(result)

Output

[1 2 3 4]

Example 3: where() with Array Condition

We can use array_like objects (such as lists, arrays etc.) as a condition in the where() method.

import numpy as np

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

# returns element of x when True 
# returns element of y when False
result = np.where([[True, True], [False, False]], x, y)

print(result)

# returns element of x when True 
# returns element of y when False
result = np.where([[True, False], [False, True]], x, y)


print(result)

Output

[[1 2]
 [-3 -4]]

[[1 -2]
 [-3 4]]

Example 4: where() with Multiple Conditions

The test condition in a where() method may have multiple conditions.

We use

• the | operator to perform OR operation on multiple conditions
• the & operator to perform AND operation on multiple conditions

import numpy as np

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

# if element is less than 2 or greater than 6, test condition is True
test_condition1 = (x < 2) | (x > 6)
# select element of x if test condition is True
# select 0 if test condition is False
result1 = np.where(test_condition1, x, 0)

print(result1)

# if element is  greater than 2 and less than 6, test condition is True
test_condition2 = (x > 2) & (x < 6)

# select element of x if test condition is True
# select 0 if test condition is False
result2 = np.where(test_condition2, x, 0)

print(result2)

Output

[1 0 0 0 0 0 7]
[0 0 3 4 5 0 0]

Example 5: where() with Only One Argument

If we pass a single argument (test condition) to numpy.where(), it tells us where in a given array the given condition is met by returning the indices.

import numpy as np

originalArray = np.array([0, 10, 20, 30, 40, 50, 60, 70])

# returns index of elements for which the test condition is True
result = np.where(originalArray > 30)

print(result)

Output

(array([4, 5, 6, 7]),)