Python Numpy.where() Function

1. Syntax of numpy.where()
2. Example Codes: numpy.where() Without [X, Y]
3. Example Codes: numpy.where() With 1-D Array
4. Example Codes: numpy.where() With 2-D Array
5. Example Codes: numpy.where() With Multiple Conditions

Numpy.where() function generates the indexes of the array that meet the input condition, if x, y are not given; or the array elements from either x or y based on the given condition.

Syntax of numpy.where()

numpy.where(condition, [x, y])

Parameters

Return

It returns an array. If condition is True, the result contains elements of x, and if condition is False, the result contains elements of y.

It returns indexes of the array is x, y are not given.

Example Codes: numpy.where() Without [X, Y]

import numpy as np

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

n = np.where(m > 3)

print(n)

Output:

(array([3, 4], dtype=int64),)

It returns the indexes of m, where its element is larger than 3 - a > 3.

If you need the element rather than the indexes.

Example Codes: numpy.where() With 1-D Array

import numpy as np

m = np.where([True, False, True], [1,2,3], [4, 5, 6])

print(m)

Output:

[1 5 3]

When the condition is a 1-D array, Numpy.where() function iterates over the condition array, and picks the element from x if the condition element is True, or the element from y if the condition element is False.

Example Codes: numpy.where() With 2-D Array

import numpy as np

x = np.array([[10, 20, 30], [3, 50, 5]])
y = np.array([[70, 80, 90], [100, 110, 120]])
condition = np.where(x > 20, x, y)

print("Input array :")
print(x)
print(y)
print("Output array with condition applied:")
print(condition)

Output:

Input array :
[[10 20 30]
[ 3 50  5]]
[[ 70  80  90]
[100 110 120]]
Output array with condition applied:
[[ 70  80  30]
[100  50 120]]

It applies the condition of x>20 to all the elements of x, if it is True, then the element of x yields as output, and if it is False, it yields the element of y.

We make a simplified example to show how it works.

import numpy as np

m = np.where(
    [[True, False, True], [False, True, False]],
    [[1, 2, 3], [4, 5, 6]],
    [[7, 8, 9], [10, 11, 12]],
)

print(m)

Output:

[[ 1  8  3]
 [10  5 12]]

Example Codes: numpy.where() With Multiple Conditions

We could also apply two or multiple conditions in the numpy.where() function.

import numpy as np

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

n = np.where((m > 1) & (m < 5), m, 0)

print(n)

Output:

[0 2 3 4 0]

It applies the multiple conditions, m > 1 and m < 5, and returns the element if the element satisfies both conditions.

The logic between the multiple conditions is not limited to AND (&), but OR (|) is also accepted.

import numpy as np

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

n = np.where((m < 2) | (m > 4), m, 0)

print(n)

Output:

[1 0 0 0 5]