Boolean masks and fancy indexing

Twelve customers rated a product from 1 to 5 stars, and counts[ratings] += 1 produced a histogram whose bars add up to 5, not 12. Indexing with an array of positions, or with a mask of True and False, is one of NumPy’s most useful tools, and this is its best-known trap: each star value was counted once, however often it appeared.

You will filter and replace values with masks, combine conditions without the errors Python’s and causes, find positions with np.nonzero, pick rows, pairs and grids with integer arrays, choose a different position in each row, and count repeats correctly with np.add.at and np.bincount.

A boolean mask selects, counts and replaces values

A comparison on an array compares every element at once and returns an array of True and False with the same shape. Used as an index, that boolean maskboolean mask: An array of True and False, usually from a comparison such as a > 50, used as an index: a[mask] gives the values where it is True as a new 1-D array, and a[mask] = v writes v there. keeps the elements where it is True. Here are eight readings from a temperature sensor, one a minute, with three glitches among them. Step through the comparisons, the mask and what it gives:

The run does the same in NumPy:

"""Find, count and replace glitches in sensor readings with a boolean mask."""
import numpy as np

readings = np.array([21.0, 85.0, 22.0, -40.0, 23.0, 24.0, 99.0, 22.0])  # degrees C, one a minute
low = readings < -10
high = readings > 50
glitch = low | high  # True where either is True
print("low:   ", low)
print("high:  ", high)
print("glitch:", glitch)
print("readings[glitch]:", readings[glitch], "| glitch.sum():", glitch.sum())
print("mean of the good readings:", readings[~glitch].mean())
print("shares memory with readings:", np.shares_memory(readings, readings[glitch]))
readings[glitch] = np.nan  # an assignment: it writes into readings
print("after readings[glitch] = np.nan:", readings)

low:    [False False False  True False False False False]
high:   [False  True False False False False  True False]
glitch: [False  True False  True False False  True False]
readings[glitch]: [ 85. -40.  99.] | glitch.sum(): 3
mean of the good readings: 22.4
shares memory with readings: False
after readings[glitch] = np.nan: [21. nan 22. nan 23. 24. nan 22.]

readings[glitch] is a new 1-D array of the three glitches: reading through a mask is advanced indexingadvanced indexing: Indexing with an array of integers or of booleans, such as a[[0, 2]] or a[a > 5]. Reading with it always makes a copy; on the left of =, it writes into the array. More in “Indexing, slicing and views”, so it copies, and the run printed shares memory with readings: False. Because True counts as 1 and False as 0, glitch.sum() counts them, and ~glitch turns the mask round, so readings[~glitch].mean() is the mean of the five good readings, 22.4. On the left of = the same mask writes into readings: the three glitches became NaN in place.

A mask of the array’s own shape always gives a 1-D result, whatever the array’s shape, because the selected values need not form rows or columns. NumPy’s guide states it for a mask with as many dimensions as the array: the result is “a 1-dimensional array”. When the shape must stay, choose values with np.where instead, two sections below.

Combine conditions with bitwise operators, in parentheses

Masks combine element by element with & (and), | (or) and ~ (not), and each comparison needs parentheses of its own. The run shows three combinations that work and the two ways it usually goes wrong:

"""Conditions combine with &, | and ~, each comparison in parentheses; and, or and missing parentheses fail."""
import numpy as np

readings = np.array([21.0, 85.0, 22.0, -40.0, 23.0, 24.0, 99.0, 22.0])
print("in 0 to 30:      ", readings[(readings >= 0) & (readings <= 30)])
print("outside 0 to 30: ", readings[(readings < 0) | (readings > 30)])
print("not above 50:    ", readings[~(readings > 50)])
try:
    readings[(readings >= 0) and (readings <= 30)]
except ValueError as error:
    print("with and:", type(error).__name__ + ":", error)
try:
    readings[readings >= 0 & readings <= 30]
except TypeError as error:
    print("without parentheses:", type(error).__name__ + ":", error)

in 0 to 30:       [21. 22. 23. 24. 22.]
outside 0 to 30:  [ 85. -40.  99.]
not above 50:     [ 21.  22. -40.  23.  24.  22.]
with and: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
without parentheses: TypeError: ufunc 'bitwise_and' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

and fails because Python’s and asks for one True or False of the array on its left, and an array of eight booleans has no single truth value: hence the ValueError that suggests a.any() or a.all(). Without parentheses the error is a different one. In Python’s table of operator precedence & binds more tightly than >= and <=, so readings >= 0 & readings <= 30 means readings >= (0 & readings) <= 30, and 0 & readings asks for a bitwise and of an integer with floats, which NumPy refuses with a TypeError. On an integer array the & would get through and the line would fail later, with the ambiguous truth value error of a chained comparison.

np.nonzero gives positions, and np.where chooses values

A mask says which elements; sometimes you need where they are. np.nonzero(mask) returns the positions of the True values, as a tuple with one array of indexes per axis. np.where does two jobs: given only a condition it is a shorthand for nonzero (its reference says to prefer nonzero then), and given three arguments, np.where(cond, x, y) builds a new array that takes x where the condition is True and y elsewhere:

"""Positions with np.nonzero, an element-wise choice with np.where, and a Python int that does not fit."""
import numpy as np

readings = np.array([21.0, 85.0, 22.0, -40.0, 23.0, 24.0, 99.0, 22.0])
glitch = (readings < -10) | (readings > 50)
print("np.nonzero(glitch):", np.nonzero(glitch))
print("np.where(glitch):  ", np.where(glitch))
cleaned = np.where(glitch, np.nan, readings)  # a new array: readings keeps its values
print("np.where(glitch, np.nan, readings):", cleaned)
print("readings[1] is still", readings[1])

grid = np.array([[3, 9, 1], [7, 2, 8]])
rows, cols = np.nonzero(grid > 5)
print("grid > 5 at rows", rows, "and columns", cols, "->", grid[rows, cols])

small = np.zeros(3, dtype=np.int8)  # int8 holds -128 to 127
try:
    np.where([True, False, True], 300, small)
except OverflowError as error:
    print("OverflowError:", error)

np.nonzero(glitch): (array([1, 3, 6]),)
np.where(glitch):   (array([1, 3, 6]),)
np.where(glitch, np.nan, readings): [21. nan 22. nan 23. 24. nan 22.]
readings[1] is still 85.0
grid > 5 at rows [0 1 1] and columns [1 0 2] -> [9 7 8]
OverflowError: Python integer 300 out of bounds for int8

The glitches are at positions 1, 3 and 6, and both calls printed the same tuple. On a 2-D array there is one array per axis: the values above 5 in grid are at rows [0 1 1] and columns [1 0 2], and those two arrays used together as an index, grid[rows, cols], give the values back, [9 7 8]. np.where(glitch, np.nan, readings) keeps all eight positions and leaves readings as it was: readings[1] is still 85.0.

An array of integers picks positions in any order

An index array lists the positions you want, in the order you want them, repeats allowed: m[[2, 0, 2]] gives rows 2, 0 and 2 again, and m[:, [3, 0]] gives every row with column 3 first. That is fancy indexingfancy indexing: Indexing with arrays of integers, such as a[[2, 0, 2]]: the result holds the values at those positions, in that order and with repeats, as a new array. One index array per axis pairs the positions up.. When you give one index array per axis, NumPy pairs them up position by position, which surprises nearly everyone the first time. Step through the same two lists used as pairs and as a grid:

The run shows the pairs, the grid and two rules about shapes:

"""Integer arrays pick positions in any order, pair up across axes, or make a grid with np.ix_."""
import numpy as np

m = np.arange(12, dtype=np.int64).reshape(3, 4)
print(m)
print("m[[2, 0, 2]]:", m[[2, 0, 2]].tolist())
print("m[:, [3, 0]]:", m[:, [3, 0]].tolist())
print("m[[0, 2], [1, 3]]:", m[[0, 2], [1, 3]].tolist())
print("m[np.ix_([0, 2], [1, 3])]:", m[np.ix_([0, 2], [1, 3])].tolist())
idx = np.array([[0, 1], [1, 2]])
print("m[idx].shape:", m[idx].shape)
x = np.zeros((2, 3, 4))
print("x[0, :, [0, 1]].shape:", x[0, :, [0, 1]].shape)
print("x[0][:, [0, 1]].shape:", x[0][:, [0, 1]].shape)

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
m[[2, 0, 2]]: [[8, 9, 10, 11], [0, 1, 2, 3], [8, 9, 10, 11]]
m[:, [3, 0]]: [[3, 0], [7, 4], [11, 8]]
m[[0, 2], [1, 3]]: [1, 11]
m[np.ix_([0, 2], [1, 3])]: [[1, 3], [9, 11]]
m[idx].shape: (2, 2, 4)
x[0, :, [0, 1]].shape: (2, 3)
x[0][:, [0, 1]].shape: (3, 2)

m[[0, 2], [1, 3]] is [m[0, 1], m[2, 3]], the two values 1 and 11. np.ix_ turns the lists into a column and a row of positions, so every listed row meets every listed column: [[1, 3], [9, 11]]. In general the result takes the shape of the index: m[idx] with a (2, 2) index gives (2, 2, 4), a row of 4 for each position of idx.

The last two lines show a rule from NumPy’s guide that catches people out. When advanced indexes are separated by a slice, as in x[0, :, [0, 1]], “the dimensions resulting from the advanced indexing operation come first”: the integer 0 and the list [0, 1] combine into one axis of length 2 at the front, so the shape is (2, 3). x[0][:, [0, 1]] indexes in two steps and gives (3, 2).

Pick a different position in each row

Sometimes every row needs its own positions: each player’s rounds sorted from worst to best, or each player’s two best rounds. np.argsort(scores, axis=1) gives each row the positions that sort it, and np.take_along_axis pairs row i of those positions with row i of the data. np.put_along_axis writes the same way:

"""Reorder each row, take each row's top two, and mark each row's best, with one index per row."""
import numpy as np

scores = np.array([[3, 9, 1, 7], [8, 2, 6, 4], [5, 5, 9, 1]])  # 3 players, 4 rounds each
order = np.argsort(scores, axis=1)  # for each row, the positions that sort it
print("argsort, axis=1:\n", order)
print("take_along_axis:\n", np.take_along_axis(scores, order, axis=1))
top2 = np.argsort(-scores, axis=1, kind="stable")[:, :2]  # largest first; a tie goes to the lower position
print("top two positions:\n", top2)
print("top two scores:\n", np.take_along_axis(scores, top2, axis=1))
best = np.argmax(scores, axis=1, keepdims=True)  # shape (3, 1)
marks = np.zeros_like(scores)
np.put_along_axis(marks, best, 1, axis=1)
print("1 at each row's best:\n", marks)
print("scores[:, top2].shape:", scores[:, top2].shape)

argsort, axis=1:
 [[2 0 3 1]
 [1 3 2 0]
 [3 0 1 2]]
take_along_axis:
 [[1 3 7 9]
 [2 4 6 8]
 [1 5 5 9]]
top two positions:
 [[1 3]
 [0 2]
 [2 0]]
top two scores:
 [[9 7]
 [8 6]
 [9 5]]
1 at each row's best:
 [[0 1 0 0]
 [1 0 0 0]
 [0 0 1 0]]
scores[:, top2].shape: (3, 3, 2)

The first player’s rounds [3, 9, 1, 7] sort through the positions [2, 0, 3, 1] into [1, 3, 7, 9]. For the top two, the run sorts -scores, so the largest come first, with kind="stable" so that equal values keep their order: the third player’s two 5s tie, and position 0 comes before position 1, giving [2, 0]. np.argmax(..., keepdims=True) keeps one position per row in a (3, 1) array, which put_along_axis uses to write a 1 at each row’s best round.

Repeated indices count once, unless you use np.add.at

Back to the ratings. counts[ratings] += 1 reads the selected values, adds 1 and writes them back, and NumPy’s guide says what that means for a repeated index: the new value “is assigned to x[1] three times, rather than being incremented 3 times”. Here is that order of events in plain Python, on four of the ratings:

counts = [0, 0, 0, 0, 0, 0]
ratings = [5, 4, 5, 3]
# 1. read every position: a copy
got = []
for r in ratings:
    got.append(counts[r])
# 2. add 1 to each value read
for k in range(len(got)):
    got[k] += 1
# 3. write back, one by one
for r, g in zip(ratings, got):
    counts[r] = g
print(counts)

[0, 0, 0, 1, 1, 1]

Both copies of position 5 were read as 0 before anything was written, so both writes stored 1. Press Edit, add another 4 to ratings and run it again: position 4 still ends at 1. np.add.at does the same addition unbuffered, one occurrence at a time, so that a position indexed more than once collects every addition. np.bincount counts the values of an integer array in one call. Here are four ways to count all twelve ratings:

"""A histogram of ratings four ways: a loop, counts[ratings] += 1, np.add.at and np.bincount."""
import numpy as np

ratings = np.array([5, 4, 5, 3, 5, 4, 2, 5, 4, 5, 1, 4])  # 12 ratings from 1 to 5 stars

loop = np.zeros(6, dtype=np.int64)
for r in ratings:
    loop[r] += 1
print("Python loop:         ", loop, "total", loop.sum())

fancy = np.zeros(6, dtype=np.int64)
fancy[ratings] += 1  # each star value is counted once, however often it appears
print("counts[ratings] += 1:", fancy, "total", fancy.sum())

at = np.zeros(6, dtype=np.int64)
np.add.at(at, ratings, 1)  # unbuffered: every occurrence adds 1
print("np.add.at:           ", at, "total", at.sum())

counts = np.bincount(ratings, minlength=6)
print("np.bincount:         ", counts, "total", counts.sum())
print("stars per value, with weights:", np.bincount(ratings, weights=ratings))
try:
    np.bincount(ratings.astype(np.float64))
except TypeError as error:
    print("np.bincount of floats: TypeError:", error)

Python loop:          [0 1 1 1 4 5] total 12
counts[ratings] += 1: [0 1 1 1 1 1] total 5
np.add.at:            [0 1 1 1 4 5] total 12
np.bincount:          [0 1 1 1 4 5] total 12
stars per value, with weights: [ 0.  1.  2.  3. 16. 25.]
np.bincount of floats: TypeError: Cannot cast array data from dtype('float64') to dtype('int64') according to the rule 'safe'

The loop, np.add.at and np.bincount agree, [0 1 1 1 4 5] with a total of 12, while counts[ratings] += 1 kept one count per star value, a total of 5. minlength=6 makes sure the result has a position for every star value, present or not, and weights adds a number per occurrence instead of 1: here the stars themselves, 16 for the four 4-star ratings and 25 for the five 5-star ones. Floats are refused with a TypeError, which the NumPy 2.5 release notes list as the end of a deprecation that began in 2.1.

Practise: count values and rank within rows

Two functions in plain Python rebuild what this lesson’s NumPy calls do: a counter that treats repeats as np.bincount does, and the top-k positions per row that argsort and take_along_axis give. Their sample tests run right here:

def count_values(values, minlength=0):
    """Return a list whose position v holds how many times v appears in values (like np.bincount)."""
    counts = [0] * max(minlength, max(values, default=-1) + 1)
    for v in set(values):
        counts[v] = 1
    return counts

from counting import count_values


def test_the_lessons_ratings():
    """counts all twelve ratings, repeats included"""
    assert count_values([5, 4, 5, 3, 5, 4, 2, 5, 4, 5, 1, 4]) == [0, 1, 1, 1, 4, 5]


def test_minlength():
    """pads the result to minlength, and grows past it when a value needs more room"""
    assert count_values([0, 0, 1, 3, 3, 3], minlength=5) == [2, 1, 0, 3, 0]
    assert count_values([2], minlength=1) == [0, 0, 1]
    assert count_values([], minlength=3) == [0, 0, 0]
    assert count_values([]) == []


def test_bad_values_are_refused():
    """refuses a negative value with ValueError and a float with TypeError"""
    try:
        count_values([1, -1])
    except ValueError:
        pass
    else:
        assert False, "a negative value must raise ValueError"
    try:
        count_values([1.0, 2.0])
    except TypeError:
        return
    assert False, "a float must raise TypeError"

def top_k_per_row(scores, k):
    """Return, for each row, the positions of its k largest values, largest first, ties by the lower position."""
    return [sorted(range(len(row)), key=lambda j: row[j])[:k] for row in scores]

from topk import top_k_per_row

PLAYERS = [[3, 9, 1, 7], [8, 2, 6, 4], [5, 5, 9, 1]]  # the lesson's scores


def test_the_lessons_top_two():
    """gives the positions that the lesson's run printed"""
    assert top_k_per_row(PLAYERS, 2) == [[1, 3], [0, 2], [2, 0]]


def test_ties_go_to_the_lower_position():
    """puts the lower column first when two values are equal"""
    assert top_k_per_row([[1, 1, 1]], 2) == [[0, 1]]
    assert top_k_per_row([[4, 7, 7, 2]], 3) == [[1, 2, 0]]


def test_k_of_zero_and_k_of_all():
    """returns empty lists for k = 0 and the whole order for k = the row length"""
    assert top_k_per_row(PLAYERS, 0) == [[], [], []]
    assert top_k_per_row(PLAYERS, 4) == [[1, 3, 0, 2], [0, 2, 3, 1], [2, 0, 1, 3]]
    assert top_k_per_row([], 2) == []


def test_k_out_of_range():
    """refuses a k below 0 or longer than a row"""
    for k in (-1, 5):
        try:
            top_k_per_row(PLAYERS, k)
        except ValueError:
            continue
        assert False, f"k = {k} must raise ValueError"

The self-check mixes all six ideas, starting with the one your answers so far show you know least well:

Interview questions

Warm-up (fresher): What does a[a > 0] return, and how is it different from np.where(a > 0, a, 0)?

a > 0 is a boolean mask, an array of True and False with the shape of a. a[a > 0] uses it as an index and returns a new 1-D array holding only the positive values, in order, so a (3, 4) array with five positive values gives shape (5,); it is a copy, so changing it leaves a alone. np.where(a > 0, a, 0) also builds a new array, but with the full shape of a: it takes the value of a where the mask is True and 0 everywhere else. Use the mask to pull values out, for example to average them, and np.where to replace values while keeping every position. To change a itself, assign through the mask: a[a < 0] = 0.

Two interview questions build on this lesson, one about a histogram that counts too little and one about an index whose shape comes out swapped; their first steps are on the module’s question page:

Key takeaways

• A boolean mask selects values into a new 1-D array, counts them with sum(), and changes them in place when it is used on the left of =.
• Conditions combine with &, | and ~, each comparison in parentheses; and, or and chained comparisons raise on arrays, and a missing parenthesis lets & bind before the comparison.
• np.nonzero returns the positions of the True values, one array per axis, while np.where(cond, x, y) builds a new array of the original shape from two choices.
• One index array per axis pairs positions up, np.ix_ makes a grid, and the result takes the shape of the index, with separated advanced indexes moving their axis to the front.
• take_along_axis and put_along_axis use a different position in each row, as argsort and argmax with keepdims=True produce them.
• x[idx] += 1 counts a repeated index once; np.add.at and np.bincount count every occurrence.