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'sand/or/not.and/orraiseValueErroron 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 toa[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.