np.where() – A Simple Illustrated Guide

Python’s numpy.where(condition, x, y) function returns an array with elements from x where condition is True, and elements from y elsewhere. 

When simply calling numpy.where(condition), it is the shorthand of np.asarray(condition).nonzero() and returns a tuple containing the indices of elements that meets the condition for each dimension.

If it sounds great to you, please continue reading, and you will fully understand the numpy.where() function through Python code snippets and vivid visualization.

This tutorial is about numpy.where() function. 

  • Concretely, I will introduce its syntax and arguments. 
  • Then, you will learn some basic examples of this function.
  • Finally, I will address three top questions about numpy.where(), including numpy.where 2d array, numpy.where multiple conditions, and numpy.where returns tuple.

You can find all the codes in this tutorial here.

Syntax and Arguments

Here is the syntax of numpy.where():

# Syntax
numpy.where(condition, [x, y, ]/)

Here is the argument table of numpy.where().

ArgumentAcceptDescription
conditionarray_like, boolWhere True, yield x, otherwise yield y.
x, yarray_like; optionalValues from which to choose.

The output of numpy.where() function depends on the argument(s) you pass. Please check out the next basic examples to have a better understanding!

Basic Example 1

Okay, letโ€™s revisit the two basic examples shown upfront.

Python’s numpy.where(condition, x, y) returns an array with elements from x where condition is True, and elements from y elsewhere. 

The code starts with importing numpy as np, creating two ndarray objects x and y, and use the condition x < 3 to call the numpy.where() method.

Then, the numpy.where(condition, x, y) does all the output array construction under the hood shown above.

Finally, we print out the output array and find ourselves understanding numpy.where a bit more!

import numpy as np

# basic example 1
x = np.arange(1, 4)
y = np.zeros(3)
print(np.where(x < 3, x, y))

Output:

Basic Example 2

Next, when simply calling numpy.where(condition), it is the shorthand of np.asarray(condition).nonzero() and returns a tuple containing the indices of elements that meets the condition for each dimension.

The code starts with importing NumPy as np, creating one ndarray object x, and use the condition x < 3 to call the numpy.where() method.

Then, the numpy.where(condition) does all the output tuple construction under the hood shown above.

Finally, we print out the output tuple and find ourselves understand numpy.where a bit more!

import numpy as np

# basic example 2
x = np.arange(1, 4)
print(np.where(x < 3))

Output:

np.where() 2d array

Letโ€™s deal with 2d array with numpy.where() in this part.

Just to refresh your memory, if calling numpy.where(condition, x, y), we can get a new ndarray with elements either from x or y depending on the condition.ย 

The code starts with importing numpy as np and creating two ndarray objects x and y (all zeros).

Then, if elements in x does not meet our condition x > 6, we will replace elements in x with the corresponding elements in y, which is zero in our case.

Finally, we print out the output array.

import numpy as np

x = np.arange(1, 13).reshape(3, 4)
y = np.zeros(12).reshape(3, 4)
# x:
# [[ 1  2  3  4]
#  [ 5  6  7  8]
#  [ 9 10 11 12]]
# y:
# [[0. 0. 0. 0.]
#  [0. 0. 0. 0.]
#  [0. 0. 0. 0.]]


new_array = np.where(x > 6, x, y)
print(new_array)
# same as:
# new_array = np.where(x > 6, x, 0)

Output:

np.where() multiple conditions

You can pass multiple conditions to numpy.where() using & for AND and | for OR logics. For example, if elements in x do not meet our multiple conditions, we will replace elements in x with the corresponding elements in y.ย 

The code starts with importing numpy as np, creating two ndarray objects x and y (all zeros), and using operators & and | respectively to create two multiple conditions.

Then, if elements in x do not meet our multiple conditions, we will replace elements in x with the corresponding elements in y, which is zero in our case.

Finally, we print out the output array.

import numpy as np

x = np.arange(12).reshape(3, 4)
y = np.zeros(12).reshape(3, 4)
# x:
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

# just to focus on the multiple conditions first!
conditions1 = (x > 3) & (x < 8)
conditions2 = (x < 3) | (x > 8)

# conditions are bool ndarray
print(f'condition 1:\n {conditions1}\n')
print(f'condition 2:\n {conditions2}\n')

# use the np.where() function
print(f'np.where(conditions1, x, y):\n {np.where(conditions1, x, y)}\n')
print(f'np.where(conditions2, x, y):\n {np.where(conditions2, x, y)}')

Output:

np.where() returns tuple

So far, we are quite familiar with the calling pattern, numpy.where(condition, x, y).

How about the (awkward) way that returns a tuple instead? Yes, I am talking about using numpy.where(condition).

Previously, I mentioned that when simply calling numpy.where(condition), it is the shorthand of np.asarray(condition).nonzero() and returns a tuple containing the indices of elements that meets the condition for each dimension.

Letโ€™s imagine we have a 2d array (two dimensions!) and want to get the row and column indices of elements meeting some conditions. We can use numpy.where(condition) to get the indices of elements meeting our condition and further our old friend zip() and list() to get the coordinate of elements!

The code starts with importing numpy as np and creating one ndarray object x.

Then, we call the numpy.where(condition) to get the indices of elements that meet the condition for each dimension (row and column in a matrix view; x-axis and y-axis in a cartesian coordinate view).

Finally, we use the zip() and list() functions to get the coordinates meeting our condition!ย 

import numpy as np

x = np.arange(12).reshape(3, 4)
# x:
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

# use the np.where() function to get the indices
a_tuple = np.where(x)
print(type(a_tuple))
print('-'*85)
print("Indices of elements that meets the condition for each dimension:")
print(np.where(x))

# use the zip() and list() function to get the coordinates
print('-'*85)
print("Coordinates (row, column) that meets the condition for each dimension:")
print(list(zip(*np.where(x))))

Output:

Summary

Thatโ€™s it for our np.where() article.ย 

We learned about its syntax, arguments, and basic examples. 

We also worked on the top three questions about the np.where() function, ranging from np.where() 2d array, np.where() multiple conditions, and np.where() returns tuple.ย 

Hope you enjoy all this and happy coding!