NumPy where()

The numpy.where() method returns a new array based on a condition applied to each element of an array.

Example

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]),)

Our premium learning platform, created with over a decade of experience and thousands of feedbacks.

Learn and improve your coding skills like never before.

Try Programiz PRO
  • Interactive Courses
  • Certificates
  • AI Help
  • 2000+ Challenges