NumPy Cheatsheet

Boolean Masking

Use this NumPy reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

Creating Boolean Masks

A boolean mask is an ndarray with dtype=bool. Any comparison on an array returns a mask.

import numpy as np

a = np.array([1, -2, 3, -4, 5, 0])

a > 0          # [True, False, True, False, True, False]
a == 0         # [False, False, False, False, False, True]
a != 3         # [True, True, False, True, True, True]
a >= 2         # [False, False, True, False, True, False]

# Multiple conditions — use & | ~ (not "and"/"or"/"not" — those work element-wise wrong)
(a > 0) & (a < 4)    # [True, False, True, False, False, False]
(a < 0) | (a == 0)   # [False, True, False, True, False, True]
~(a > 0)             # [False, True, False, True, False, True]

Gotcha: use &, |, ~ with parentheses around each condition, not Python's and/or/not. and/or raise ValueError on arrays.

Applying Masks (Fancy Indexing)

a = np.array([10, -20, 30, -40, 50])
mask = a > 0

a[mask]          # [10, 30, 50]   — always returns a 1-D copy
a[a > 0]         # same, inline

# 2-D masking (returns flattened 1-D result)
b = np.array([[1, -2], [3, -4]])
b[b > 0]         # [1, 3]

# Conditional select (doesn't flatten)
np.where(mask, a, 0)         # [10, 0, 30, 0, 50]  — keep positives, else 0
np.where(a > 0, a, -a)       # abs(a) equivalent

In-place Assignment via Mask

a = np.array([1.0, -2.0, 3.0, -4.0, 5.0])
a[a < 0] = 0         # zero-out negatives
a[a > 3] *= 2        # double values above 3

# Reassign to a value array (must match count or broadcast)
a[a > 0] = np.array([10, 30, 50])   # one value per True element (count must match)
a[a > 0] = 99                        # scalar broadcasts to all True positions

np.where

np.where(condition)             # equivalent to np.nonzero — returns tuple of index arrays
np.where(condition, x, y)       # element-wise choose: x where True, y where False

a = np.arange(-3, 4)
np.where(a >= 0, a, -a)        # abs — same as np.abs(a)
np.where(a > 0, "pos", "neg")  # works with any broadcastable x, y

# Multi-condition with np.select
conditions = [a < 0, a == 0, a > 0]
choices    = [-1,    0,       1    ]
np.select(conditions, choices, default=np.nan)   # like if/elif/else

np.nonzero and np.argwhere

a = np.array([0, 3, 0, 5, 0, 1])
np.nonzero(a)          # (array([1, 3, 5]),)  — tuple of arrays (one per axis)
np.nonzero(a)[0]       # [1, 3, 5]

b = np.array([[0, 1], [2, 0]])
np.nonzero(b)          # (array([0, 1]), array([1, 0]))  — row and col indices
np.argwhere(b)         # [[0,1],[1,0]]  — each row is [row, col] of a nonzero

np.flatnonzero(a)      # flat indices of nonzero: [1, 3, 5]
np.count_nonzero(a)    # 3
np.count_nonzero(b, axis=0)  # per-column nonzero count

Masked Array (numpy.ma)

numpy.ma provides arrays where some elements are masked (excluded from computations).

import numpy.ma as ma

data = np.array([1.0, -999.0, 3.0, -999.0, 5.0])
masked = ma.masked_equal(data, -999.0)    # mask where equal to fill value
masked = ma.masked_where(data < 0, data)  # mask where condition is True
masked = ma.masked_invalid(data)          # mask NaN and inf
masked = ma.masked_outside(data, 0, 10)  # mask outside range
masked = ma.masked_inside(data, -1, 1)   # mask inside range

masked.data          # underlying raw array (including masked values)
masked.mask          # boolean mask array
masked.filled(0.0)   # return ndarray with masked values replaced by fill
masked.compressed()  # 1-D array of unmasked values only

# Arithmetic ignores masked values
ma.array([1, 2, 3], mask=[0, 1, 0]).mean()   # (1 + 3) / 2 = 2.0

Creating masked arrays

ma.array([1, 2, 3], mask=[False, True, False])  # mask middle element
ma.array([1, 2, 3], mask=True)                   # mask all
ma.array([1, 2, 3], mask=False)                  # mask none

# From boolean mask
raw = np.array([1.0, np.nan, 3.0])
masked = ma.masked_invalid(raw)

Operations on masked arrays

a = ma.array([1, 2, 3, 4], mask=[0, 1, 0, 1])
a.sum()      # 4    (ignores 2 and 4)
a.mean()     # 2.0  (mean of 1 and 3)
a.min()      # 1
a + 10       # ma.array([11, --, 13, --])

Combining Conditions

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

# Multiple comparisons
(a > 2) & (a < 5)         # [F, F, T, T, F, F]
(a % 2 == 0) | (a > 4)   # [F, T, F, T, T, T]
np.logical_and(a > 2, a < 5)  # same as &, works with scalars too

# Check element membership
np.isin(a, [2, 4, 6])           # [F, T, F, T, F, T]
np.isin(a, [2, 4, 6], invert=True)  # complement

# Any / All over mask
mask = a > 3
mask.any()       # True
mask.all()       # False
np.any(mask)     # same
np.all(mask, axis=0)  # axis-aware

Boolean Mask as Integer

# Python/NumPy: True == 1, False == 0
mask = np.array([True, False, True, True])
mask.sum()      # 3  — count of True
mask.mean()     # 0.75  — fraction of True
mask.astype(int)  # [1, 0, 1, 1]

Advanced: np.extract and np.place

a = np.arange(10)
np.extract(a % 3 == 0, a)        # [0, 3, 6, 9]  — like a[a%3==0]

# np.place — put values at mask positions (modifies in-place)
np.place(a, a % 3 == 0, [99, 88, 77, 66])
# a becomes [99, 1, 2, 88, 4, 5, 77, 7, 8, 66]

Performance Notes

  • Boolean masking returns a copy, never a view.
  • For large arrays, np.where(condition, x, y) avoids Python-level iteration.
  • np.compress(mask, a) is equivalent to a[mask] but slightly lower overhead for 1-D.
  • np.take(a, np.where(condition)[0]) can be faster than fancy indexing in some cases due to contiguous access patterns.