Filtering Arrays

Filtering keeps the values that meet a condition. In NumPy, the condition is itself an array of True and False, called a mask. The filter runs in compiled code, so there is no Python loop to write.

Boolean masks

A comparison gives a mask of the same shape. Put the mask inside the brackets to keep only the True positions. True counts as 1, so mask.sum() counts the matches.

import numpy as np

scores = np.array([45, 82, 67, 91, 38, 74])
mask = scores >= 70
print(mask)
print(scores[mask])
print(mask.sum())

Output:

[False  True False  True False  True]
[82 91 74]
3

Combine conditions

Use & for and, | for or, and ~ for not. The words and and or do not work on arrays. Put each comparison in parentheses, because & and | bind tighter than > and <.

import numpy as np

scores = np.array([45, 82, 67, 91, 38, 74])
print(scores[(scores >= 50) & (scores < 80)])
print(scores[(scores < 40) | (scores > 90)])
print(scores[~(scores >= 70)])

Output:

[67 74]
[91 38]
[45 67 38]

np.where with one argument

np.where(condition) returns the positions where the condition is True, as a tuple with one array per dimension.

import numpy as np

scores = np.array([45, 82, 67, 91, 38, 74])
spots = np.where(scores >= 70)
print(spots)
print(scores[spots])

Output:

(array([1, 3, 5]),)
[82 91 74]

np.where with three arguments

np.where(condition, x, y) keeps the shape. It takes x where the condition is True and y where it is False. x and y can be single values or arrays of the same shape.

import numpy as np

scores = np.array([45, 82, 67, 91, 38, 74])
print(np.where(scores >= 50, "pass", "retry"))
print(np.where(scores < 50, 50, scores))

bonus = scores + 10
print(np.where(scores >= 70, bonus, scores))

Output:

['retry' 'pass' 'pass' 'pass' 'retry' 'pass']
[50 82 67 91 50 74]
[ 45  92  67 101  38  84]

np.extract

np.extract(condition, arr) returns only the matching values, like arr[condition]. The result is shorter than the original, so it uses less memory than a full-size where result.

import numpy as np

ages = np.array([15, 99, 87, 21, 67])
print(np.extract(ages < 90, ages))
print(np.extract((ages < 20) | (ages > 80), ages))

Output:

[15 87 21 67]
[15 99 87]

Missing values

In a float array, a missing value is np.nan. nan is not equal to anything, not even itself, so data == np.nan is always False. Use np.isnan to find it.

import numpy as np

data = np.array([4.0, np.nan, 7.5, np.nan, 2.0])
print(data == np.nan)
print(np.isnan(data))
print(np.where(np.isnan(data), 0, data))
print(data[~np.isnan(data)])

Output:

[False False False False False]
[False  True False  True False]
[4.  0.  7.5 0.  2. ]
[4.  7.5 2. ]

Key takeaways

• A comparison gives a mask. arr[mask] keeps the True values

• Use &, |, and ~, with parentheses around each comparison

• np.where(cond) gives positions. np.where(cond, x, y) keeps the shape

• Find missing floats with np.isnan, never == np.nan

• Next: Append, Insert, and Delete

• Previous: Sorting and Comparing