numpy Filtering data Filtering data with a boolean array

Help us to keep this website almost Ad Free! It takes only 10 seconds of your time:
> Step 1: Go view our video on YouTube: EF Core Bulk Extensions
> Step 2: And Like the video. BONUS: You can also share it!

Example

When only a single argument is supplied to numpy's where function it returns the indices of the input array (the condition) that evaluate as true (same behaviour as numpy.nonzero). This can be used to extract the indices of an array that satisfy a given condition.

import numpy as np

a = np.arange(20).reshape(2,10)
# a = array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
#           [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

# Generate boolean array indicating which values in a are both greater than 7 and less than 13
condition = np.bitwise_and(a>7, a<13)
# condition = array([[False, False, False, False, False, False, False, False,  True, True],
#                    [True,  True,  True, False, False, False, False, False, False, False]], dtype=bool)

# Get the indices of a where the condition is True
ind = np.where(condition)
# ind = (array([0, 0, 1, 1, 1]), array([8, 9, 0, 1, 2]))

keep = a[ind]
# keep = [ 8  9 10 11 12]

If you do not need the indices, this can be achieved in one step using extract, where you agian specify the condition as the first argument, but give the array to return the values from where the condition is true as the second argument.

# np.extract(condition, array)
keep = np.extract(condition, a)
# keep = [ 8  9 10 11 12]

Two further arguments x and y can be supplied to where, in which case the output will contain the values of x where the condition is True and the values of y where the condition is False.

# Set elements of a which are NOT greater than 7 and less than 13 to zero, np.where(condition, x, y)
a = np.where(condition, a, a*0)
print(a)
# Out: array([[ 0,  0,  0,  0,  0,  0,  0,  0,  8,  9],
#            [10, 11, 12,  0,  0,  0,  0,  0,  0,  0]])


Got any numpy Question?