NumPy
N-dimensional arrays, dtypes, indexing, broadcasting, vectorized math, linear algebra and random numbers with NumPy 2.5 on Python 3.14. Labeled tables are in pandas.
Creating arrays
uv add numpy # pip install numpy
uv run --with numpy python # throwaway REPL| Call | Gives |
|---|---|
np.array([[1, 2], [3, 4]]) | from nested sequences; dtype inferred |
np.asarray(x) | same, but no copy if x is already a matching ndarray |
np.zeros((2, 3)), np.ones(4) | filled with 0.0 / 1.0 (float64) |
np.full((2, 2), 7) | filled with a value; dtype from the value |
np.empty(1000) | uninitialised memory (fast; fill it before reading) |
np.zeros_like(a) | same shape and dtype as a; also ones_like, full_like, empty_like |
np.arange(0, 10, 2) | [0 2 4 6 8]; stop is exclusive |
np.linspace(0, 1, 5) | 5 evenly spaced points; stop is inclusive |
np.logspace(0, 3, 4) | [1 10 100 1000] |
np.eye(3), np.identity(3) | identity matrix (eye takes k= for an off-diagonal) |
np.diag([1, 2, 3]) | diagonal matrix; on a 2-D array it extracts the diagonal |
np.fromfunction(f, (3, 3)) | calls f(i, j) once with whole index grids |
np.fromiter(gen, dtype=float) | from any iterator |
np.meshgrid(x, y) | coordinate grids for evaluating f(x, y) |
np.tile(a, 2), np.repeat(a, 2) | repeat the whole array / each element |
import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6]])
a.dtype, a.shape # (dtype('int64'), (2, 3))
np.arange(0, 1, 0.25) # [0. 0.25 0.5 0.75]
np.linspace(0, 1, 5) # [0. 0.25 0.5 0.75 1. ]
np.fromfunction(lambda i, j: 10 * i + j, (2, 3), dtype=int)
# [[ 0 1 2]
# [10 11 12]]Prefer linspace over a float arange step: rounding can add or drop the last point.
Dtypes & casting
| dtype | Code | Notes |
|---|---|---|
bool | ? | 1 byte per element |
int8 int16 int32 int64 | i1 i2 i4 i8 | int64 is the default integer (Windows too since 2.0) |
uint8 … uint64 | u1 … u8 | image pixels are usually uint8 |
float16 float32 float64 | f2 f4 f8 | float64 is the default float |
complex64 complex128 | c8 c16 | |
str_ | U10 | fixed-width Unicode; padded to the longest string |
StringDType() | T | variable-width UTF-8 strings (2.0+) |
bytes_ | S10 | fixed-width bytes |
datetime64[D], timedelta64[s] | M8, m8 | unit in brackets: Y M W D h m s ms us ns |
object | O | pointers to Python objects; slow, avoid |
x = np.array([1.7, -2.5, 300.0])
x.astype(np.int64) # [ 1 -2 300] truncates
np.rint(x).astype(np.int64) # [ 2 -2 300] round first
x.astype(np.uint8) # [ 1 254 44] silent wrap
x.astype(np.float32, copy=False) # no copy if same dtype
# 2.4+: refuse lossy casts instead of wrapping
x.astype(np.uint8, casting="same_value") # ValueError
np.can_cast(np.int64, np.float64) # True
np.result_type(np.int8, np.float32) # float32
np.iinfo(np.int32).max # 2147483647
np.finfo(np.float64).eps # 2.220446049250313e-16Python scalars are "weak" (NEP 50): they adopt the array's dtype instead of upcasting it.
| Expression | Result dtype |
|---|---|
np.array([1.0], np.float32) * 2.5 | float32 (Python float is weak) |
np.array([1], np.int8) + np.int64(1) | int64 (NumPy scalar is strong) |
np.array([1], np.uint8) + 300 | OverflowError: 300 does not fit in uint8 |
np.array([1, 2]) / 2 | float64: true division always gives a float |
np.array([1, 2]) // 2 | int64 |
Shape & reshaping
| Attribute / call | Meaning |
|---|---|
a.shape, a.ndim, a.size | dims tuple, number of dims, element count |
a.dtype.itemsize, a.nbytes | bytes per element, total bytes |
a.reshape(3, -1) | new shape; -1 is inferred; view when possible |
a.ravel() | flatten to 1-D, a view when possible |
a.flatten() | flatten to 1-D, always a copy |
a.T, a.transpose(2, 0, 1) | reverse / permute axes (view) |
np.swapaxes(a, 0, 1), np.moveaxis(a, 0, -1) | swap two axes / move one (view) |
a[:, np.newaxis], a[:, None] | insert a length-1 axis |
np.expand_dims(a, 0), np.squeeze(a) | add / drop length-1 axes |
v = np.arange(6) # [0 1 2 3 4 5]
m = v.reshape(2, 3) # [[0 1 2] [3 4 5]]
m.T.shape # (3, 2)
v[:, None].shape # (6, 1) column vector
v[None, :].shape # (1, 6) row vector
np.arange(24).reshape(2, 3, 4).transpose(2, 0, 1).shape
# (4, 2, 3)Setting a.shape = ... is deprecated (2.5); call reshape.
Stacking & splitting
| Call | Joins / splits along |
|---|---|
np.concatenate([a, b], axis=0) | an existing axis |
np.stack([a, b], axis=0) | a new axis (all shapes equal) |
np.vstack, np.hstack | rows / columns (1-D inputs become rows) |
np.column_stack([x, y]) | 1-D arrays as columns |
np.block([[A, B], [C, D]]) | block matrix |
np.split(a, 3), np.split(a, [2, 5]) | equal parts / at indices |
np.array_split(a, 3) | like split, uneven parts allowed |
np.vsplit, np.hsplit | rows / columns |
x, y = np.array([1, 2]), np.array([3, 4])
np.stack([x, y]).shape # (2, 2) new axis 0
np.stack([x, y], axis=1) # [[1 3] [2 4]]
np.concatenate([x, y]) # [1 2 3 4]
np.column_stack([x, y]).shape # (2, 2)
np.split(np.arange(6), [2, 5])
# [array([0, 1]), array([2, 3, 4]), array([5])]np.row_stack was removed in 2.5; use np.vstack. Growing an array in a loop copies
every time: collect a list and stack once.
Indexing & views
| Syntax | Kind | Returns |
|---|---|---|
a[2], a[1, 2] | integer | element (or sub-array of a lower dim) |
a[1:4], a[::2], a[::-1] | basic slice | view |
a[:, 0], a[..., 0] | basic slice (... fills remaining axes) | view |
a[[0, 2]] | fancy (integer array) | copy |
a[a > 0] | boolean mask | copy, always 1-D |
a[np.ix_(rows, cols)] | fancy on each axis | copy of the sub-grid |
a[[0, 1], [2, 3]] | paired fancy | elements (0, 2) and (1, 3) |
a = np.arange(12).reshape(3, 4)
a[1, 2], a[1][2] # both 6; a[1, 2] is faster
a[:, 1] # [1 5 9] column 1
a[1:, ::2] # [[ 4 6] [ 8 10]]
a[[0, 2]] # rows 0 and 2
a[[0, 2], [1, 3]] # [ 1 11]
a[np.ix_([0, 2], [1, 3])] # [[ 1 3] [ 9 11]]
a[a % 5 == 0] # [ 0 5 10]
a[(a > 2) & (a < 6)] # [3 4 5] needs parenthesesAssigning through any index writes into a (a[a < 0] = 0), but a fancy or boolean
index read gives a copy, so a[[0, 2]][0] = 99 changes nothing.
Views vs copies
| Makes a view | Makes a copy |
|---|---|
basic slicing, a.T, transpose, swapaxes | fancy and boolean indexing |
reshape/ravel when memory allows | flatten, a.copy(), np.array(a) |
a.view(np.int32) (reinterpret bytes) | astype (unless copy=False and same dtype) |
np.asarray(a) on an ndarray | arithmetic results (a + 1) |
a = np.arange(5)
s = a[1:4]
s[0] = 100
a # [ 0 100 2 3 4]
s.base is a # True
np.shares_memory(a, a[[1, 2]]) # False
safe = a[1:4].copy() # detach before mutatingA small slice keeps a huge base array alive; .copy() it if the original can go.
Broadcasting
Operations between shapes that differ follow three rules, applied from the last axis backwards:
- If the ranks differ, pad the shorter shape with 1s on the left.
- Two dims are compatible if they are equal or one of them is 1.
- A dim of 1 is stretched (without copying) to match the other.
x[:, None] shape (3, 1) [[0] y [10 20 30 40]
y shape (4,) [1]
[2]]
step 1: pad y (3, 1) and (1, 4)
step 2: compare 3 vs 1 -> 3, 1 vs 4 -> 4
result (3, 4)
x+y = [[10 20 30 40] x is copied across columns,
[11 21 31 41] y is copied down rows
[12 22 32 42]]x = np.arange(3)
y = np.array([10, 20, 30, 40])
(x[:, None] + y).shape # (3, 4) outer sum
m = np.ones((4, 3))
(m - m.mean(axis=0)).shape # (4, 3): (4,3) with (3,)
np.broadcast_shapes((5, 1, 3), (4, 1)) # (5, 4, 3)
x + y # ValueError: could not be broadcast (3,) (4,)| Shapes | Result |
|---|---|
(4, 3) and (3,) | (4, 3): row added to every row |
(4, 3) and (4,) | error; use (4, 1) via v[:, None] |
(4, 3) and (4, 1) | (4, 3): column added to every column |
(8, 1, 6) and (7, 1) | (8, 7, 6) |
(3,) and () scalar | (3,) |
Ufuncs & vectorization
Universal functions run a compiled loop element by element and broadcast their inputs.
| Group | Functions / operators |
|---|---|
| Arithmetic | + - * / // % **, np.add, np.divide, np.power, np.negative |
| Math | np.abs, np.sqrt, np.exp, np.log, np.log1p, np.sin, np.hypot |
| Rounding | np.round(a, 2), np.floor, np.ceil, np.trunc (np.fix is deprecated) |
| Element-wise max/min | np.maximum(a, b), np.minimum, np.clip(a, 0, 1) |
| Comparison | == != < >, np.isclose, np.allclose, np.array_equal |
| Logical | &, |, ~, ^ on bool arrays; np.logical_and |
| Choice | np.where(cond, a, b), np.select(conds, choices, default) |
| Ufunc option / method | Does |
|---|---|
out=buf | writes into an existing array (no allocation) |
where=mask | computes only where mask is true (pair with out=) |
dtype=np.float32 | computes in that dtype |
np.add.reduce(a) | fold along an axis (np.sum uses it) |
np.add.accumulate(a) | running result (cumsum) |
np.multiply.outer(x, y) | every pair, shape x.shape + y.shape |
np.add.at(a, idx, 1) | unbuffered in-place update; repeated indices all count |
t = np.array([-1.5, 0.2, 3.7])
np.where(t > 0, t, 0) # [0. 0.2 3.7]
np.clip(t, 0, 1) # [0. 0.2 1. ]
np.select([t < 0, t < 1], ["neg", "small"], "big")
# ['neg' 'small' 'big']
np.multiply.outer([1, 2], [1, 10, 100])
# [[ 1 10 100]
# [ 2 20 200]]
counts = np.zeros(3, dtype=int)
np.add.at(counts, [0, 0, 2], 1)
counts # [2 0 1]np.vectorize is a Python loop with broadcasting; it adds convenience, not speed.
Reductions & axis
axis names the dimension that disappears.
a = [[1, 2, 3], a.sum() -> 21
[4, 5, 6]] a.sum(axis=0) -> [5 7 9] down rows
a.sum(axis=1) -> [ 6 15] across cols
shape (2, 3) axis=0 -> (3,) axis=1 -> (2,)| Function | Notes |
|---|---|
sum, prod, cumsum, cumprod | cumsum keeps the shape |
mean, median, average(a, weights=w) | average is the weighted mean |
std, var | population by default; ddof=1 for sample |
min, max, ptp | ptp is max minus min |
argmin, argmax | index of the first extreme (flat unless axis=) |
percentile(a, 90), quantile(a, 0.9) | method="linear" by default |
any, all, count_nonzero | booleans |
np.unique(a, return_counts=True) | sorted distinct values and counts |
a = np.array([[1, 2, 3], [4, 5, 6]])
a.sum(axis=0) # [5 7 9]
a.mean(axis=1) # [2. 5.]
a.max(axis=1, keepdims=True) # [[3] [6]] shape (2, 1)
a / a.sum(axis=1, keepdims=True) # rows sum to 1
a.sum(axis=(0, 1)) # 21
np.std([2, 4, 4, 4, 5, 5, 7, 9]) # 2.0
np.std([2, 4, 4, 4, 5, 5, 7, 9], ddof=1)
# 2.138089935299395keepdims=True leaves a length-1 axis, so the result broadcasts back against the input.
NaN handling
nan poisons ordinary reductions, and nan != nan, so test with np.isnan.
| Need | Use |
|---|---|
| Detect | np.isnan(a), np.isfinite(a), np.isinf(a) |
| Count | np.isnan(a).sum() |
| Drop | a[~np.isnan(a)] |
| Replace | np.nan_to_num(a, nan=0.0), np.where(np.isnan(a), fill, a) |
| Reduce, ignoring NaN | np.nansum, nanmean, nanstd, nanmin, nanmax, nanmedian, nanpercentile, nanargmax |
| Compare arrays with NaN | np.array_equal(a, b, equal_nan=True), np.isclose(..., equal_nan=True) |
| Rows without NaN | m[~np.isnan(m).any(axis=1)] |
a = np.array([1.0, np.nan, 3.0])
a.mean(), np.nanmean(a) # (nan, 2.0)
np.nan == np.nan # False
np.isnan(a) # [False True False]
np.nan_to_num(a, nan=-1) # [ 1. -1. 3.]
np.array([1, 2])[0] = np.nan # ValueError: int has no NaNIntegers can't hold NaN. Use a float array, a sentinel, or np.ma.masked_invalid(a)
for a masked array.
Sorting & searching
| Call | Gives |
|---|---|
np.sort(a) / a.sort() | sorted copy / sort in place (last axis by default) |
np.sort(a, descending=True) | descending (2.5+); older: np.sort(a)[::-1] |
np.argsort(a) | indices that would sort a (stable=True for a stable sort) |
np.lexsort((b, a)) | sort by a, then b (last key is primary) |
np.argmax(a), np.argmin(a) | index of the max / min |
np.partition(a, k), np.argpartition | the k-th element in its sorted place, smaller ones before, O(n) |
np.where(cond), np.nonzero(a) | tuple of index arrays where true / non-zero |
np.argwhere(cond) | (n, ndim) array of coordinates |
np.flatnonzero(cond) | flat indices |
np.searchsorted(sorted_a, v) | insertion point that keeps sorted_a sorted |
np.isin(a, values) | membership mask (np.in1d was removed in 2.4) |
np.unravel_index(i, a.shape) | flat index to coordinates |
a = np.array([30, 10, 50, 20])
np.argsort(a) # [1 3 0 2]
a[np.argsort(a)[::-1]] # [50 30 20 10]
np.sort(a, descending=True) # [50 30 20 10]
np.argmax(a) # 2
np.where(a > 15) # (array([0, 2, 3]),)
np.searchsorted([10, 20, 30], [5, 20, 35]) # [0 1 3]
m = np.array([[3, 9], [7, 1]])
np.unravel_index(m.argmax(), m.shape) # (0, 1) as np ints
np.argsort(m, axis=1) # [[0 1] [1 0]] per row
top2 = np.argpartition(a, -2)[-2:] # unordered top-2 idx
np.sort(a[top2]) # [30 50]Linear algebra
| Call | Does |
|---|---|
A @ B, np.matmul(A, B) | matrix product; stacks of matrices broadcast |
np.dot(a, b), np.vecdot(x, y) | dot product / batched vector dot (2.0+) |
np.outer(x, y), np.cross(x, y) | outer product, 3-D cross product |
np.einsum("ij,jk->ik", A, B) | index-notation products (optimize=True) |
np.linalg.solve(A, b) | solves A @ x = b (better than inv(A) @ b) |
np.linalg.inv, np.linalg.pinv | inverse / pseudo-inverse |
np.linalg.det, np.linalg.matrix_rank | determinant, rank |
np.linalg.eig(A) | eigenvalues and vectors; always complex (2.5) |
np.linalg.eigh(A) | symmetric or Hermitian: real, ascending eigenvalues |
np.linalg.svd(A, full_matrices=False) | U, S, Vh with A = U @ diag(S) @ Vh |
np.linalg.norm(x), norm(A, axis=1) | L2 norm (also ord=1, np.inf, "fro") |
np.linalg.lstsq(A, b) | least squares fit |
np.linalg.qr, np.linalg.cholesky | decompositions |
A = np.array([[3.0, 1.0], [1.0, 2.0]])
b = np.array([9.0, 8.0])
x = np.linalg.solve(A, b) # [2. 3.]
np.allclose(A @ x, b) # True
np.linalg.eigh(A).eigenvalues # [1.38196601 3.61803399]
U, S, Vh = np.linalg.svd(A)
np.allclose(U * S @ Vh, A) # True
np.linalg.norm([3, 4]) # 5.0
np.einsum("ij,ij->i", A, A) # [10. 5.] row dots* is element-wise, not matrix multiplication. For sparse matrices or more routines
(expm, lu), reach for scipy.linalg and scipy.sparse.
Random numbers
Create one Generator and pass it around. The legacy np.random.seed and np.random.rand
share global state and are kept only for old code.
rng = np.random.default_rng(42) # seeded: reproducible
rng.integers(0, 10, size=5) # [0 7 6 4 4]
rng.random(3) # floats in [0, 1)
rng.normal(loc=0, scale=1, size=(2, 3))
rng.choice(["a", "b", "c"], size=4, p=[0.5, 0.3, 0.2])
rng.choice(10, size=3, replace=False) # sample, no repeats
rng.permutation(5) # shuffled arange(5)
children = rng.spawn(4) # independent streams| Legacy | Modern (rng = np.random.default_rng()) |
|---|---|
np.random.seed(0) | np.random.default_rng(0) |
np.random.rand(3) | rng.random(3) |
np.random.randn(3) | rng.standard_normal(3) |
np.random.randint(0, 10) | rng.integers(0, 10) (endpoint=True to include 10) |
np.random.shuffle(a) | rng.shuffle(a) in place / rng.permuted(a, axis=1) |
np.random.normal(...) | rng.normal(...), also uniform, poisson, binomial, exponential |
Streams differ between the two APIs; seeding the legacy one doesn't affect a Generator.
Use rng.spawn(n) (or SeedSequence.spawn) for workers, never seed + i.
Saving & loading
| Call | Format |
|---|---|
np.save("a.npy", a) / np.load("a.npy") | one array, binary, keeps dtype and shape |
np.savez("d.npz", x=x, y=y) | several arrays in a zip; savez_compressed for zlib |
np.load("d.npz") | lazy NpzFile; index by name, use as a context manager |
np.load(p, mmap_mode="r") | memory-map a .npy instead of reading it all |
np.savetxt("a.csv", a, delimiter=",", fmt="%.3f") | text |
np.loadtxt(p, delimiter=",", skiprows=1) | clean numeric text |
np.genfromtxt(p, delimiter=",", missing_values="") | text with gaps (fills nan) |
a.tofile(p) / np.fromfile(p, dtype) | raw bytes, no header (you track dtype and shape) |
from pathlib import Path
from tempfile import mkdtemp
d = Path(mkdtemp())
x, y = np.arange(3), np.eye(2)
np.save(d / "x.npy", x)
np.savez_compressed(d / "xy.npz", x=x, y=y)
np.load(d / "x.npy") # [0 1 2]
with np.load(d / "xy.npz") as data:
sorted(data.files) # ['x', 'y']
data["y"].shape # (2, 2)np.load has allow_pickle=False by default; don't enable it for untrusted files, since
unpickling runs code. For tabular or cross-language data use Parquet via pandas or pyarrow.
Performance
| Tip | Why |
|---|---|
| Replace Python loops with array ops | per-element work moves into C/SIMD; often 10-100x faster |
Preallocate, or collect a list then np.array once | np.append and concatenate in a loop copy each time |
In-place ops: a += 1, np.multiply(a, 2, out=a) | no temporary arrays |
Reuse buffers with out= | avoids allocation inside hot loops |
| Pick the smallest safe dtype | float32 halves memory and bandwidth |
| Reduce along the contiguous axis | C order (default) means the last axis is contiguous |
np.ascontiguousarray(a) before heavy work | a transposed or strided view is slower to traverse |
| Slice (a view) instead of fancy indexing | fancy indexing copies |
np.einsum(..., optimize=True) | picks a good contraction order |
np.lib.stride_tricks.sliding_window_view | windows with no copy |
Numba (@njit) or Cython | when a loop genuinely can't be vectorized |
a = np.ones((1000, 1000))
a.flags["C_CONTIGUOUS"], a.T.flags["C_CONTIGUOUS"]
# (True, False)
buf = np.empty_like(a)
np.multiply(a, 2.0, out=buf) # no new allocation
np.add(buf, 1.0, out=buf) # in place
a.astype(np.float32).nbytes # 4000000 (half of 8e6)Time with python -m timeit or %timeit in IPython; np.show_config() shows which
BLAS the linear algebra uses.
Typing
from typing import Any
import numpy as np
import numpy.typing as npt
type F64 = npt.NDArray[np.float64]
def zscore(x: npt.ArrayLike) -> F64:
arr = np.asarray(x, dtype=np.float64)
return (arr - arr.mean()) / arr.std()
def cast(a: F64, dtype: npt.DTypeLike) -> npt.NDArray[Any]:
return a.astype(dtype)
# shape-typed: ndarray[shape, dtype]
Matrix = np.ndarray[tuple[int, int], np.dtype[np.float64]]| Type | Use for |
|---|---|
npt.NDArray[np.float64] | an ndarray with that dtype, any shape |
npt.NDArray[np.floating] | any float dtype |
npt.ArrayLike | anything np.asarray accepts (lists, scalars, arrays) |
npt.DTypeLike | anything np.dtype() accepts ("f8", np.int32, float) |
np.ndarray[tuple[int, int], np.dtype[np.float64]] | 2-D float array (shape typing, improving in 2.5) |
NumPy ships its own stubs; run uvx mypy or uvx pyright with numpy installed in the
environment. Annotate public boundaries with ArrayLike in and NDArray out.
Recipes
Normalize columns
Scale each feature to zero mean and unit variance, or to the range 0 to 1.
X = np.array([[1.0, 200.0], [2.0, 300.0], [3.0, 400.0]])
mu = X.mean(axis=0)
sd = X.std(axis=0)
Z = (X - mu) / np.where(sd == 0, 1, sd) # no div by 0
Z.round(3)
# [[-1.225 -1.225]
# [ 0. 0. ]
# [ 1.225 1.225]]
lo, hi = X.min(axis=0), X.max(axis=0)
(X - lo) / (hi - lo) # each column in [0, 1]Pairwise distances via broadcasting
Distance from every point in P to every point in Q without a loop.
P = np.array([[0.0, 0.0], [3.0, 4.0]]) # (n, d)
Q = np.array([[0.0, 0.0], [6.0, 8.0], [3.0, 0.0]]) # (m, d)
diff = P[:, None, :] - Q[None, :, :] # (n, m, d)
D = np.sqrt((diff**2).sum(axis=-1)) # (n, m)
D
# [[ 0. 10. 3.]
# [ 5. 5. 4.]]
D.argmin(axis=1) # nearest Q: [0 2]For large inputs this needs n * m * d memory; scipy.spatial.distance.cdist avoids it.
Moving average
Smooth a 1-D signal over a window of k points.
from numpy.lib.stride_tricks import sliding_window_view
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
k = 3
np.convolve(x, np.ones(k) / k, mode="valid") # [2. 3. 4. 5.]
sliding_window_view(x, k).mean(axis=1) # same, no copy
c = np.cumsum(np.insert(x, 0, 0.0)) # O(n)
(c[k:] - c[:-k]) / k # [2. 3. 4. 5.]One-hot encode
Turn integer class labels into indicator columns.
labels = np.array([2, 0, 1, 2])
n = labels.max() + 1
np.eye(n, dtype=np.int8)[labels]
# [[0 0 1]
# [1 0 0]
# [0 1 0]
# [0 0 1]]
cats, codes = np.unique(["b", "a", "b"], return_inverse=True)
cats, codes # ['a' 'b'], [1 0 1]Find rows matching a condition
Filter a 2-D array by column values and get the row numbers.
data = np.array([[1, 25, 0], [2, 40, 1], [3, 31, 1]])
mask = (data[:, 1] > 30) & (data[:, 2] == 1)
data[mask] # [[ 2 40 1] [ 3 31 1]]
np.flatnonzero(mask) # [1 2]
target = np.array([2, 40, 1])
np.flatnonzero((data == target).all(axis=1)) # [1]
np.isin(data[:, 0], [1, 3]) # [ True False True]Bin values
Assign each value to a bucket, or count values per bucket.
ages = np.array([3, 17, 18, 42, 65, 90])
edges = np.array([0, 18, 65])
np.digitize(ages, edges) # [1 1 2 2 3 3]
labels = np.array(["child", "adult", "senior"])
labels[np.digitize(ages, edges) - 1]
# ['child' 'child' 'adult' 'adult' 'senior' 'senior']
counts, bin_edges = np.histogram(ages, bins=[0, 18, 65, 120])
counts # [2 2 2]
np.bincount(np.digitize(ages, edges)) # [0 2 2 2]digitize puts a value equal to an edge in the bin to its right (right=True flips that).
Top-k per row
Get the k largest values in each row, in order, in O(n) per row.
S = np.array([[5, 1, 9, 3], [2, 8, 4, 7]])
k = 2
idx = np.argpartition(S, -k, axis=1)[:, -k:]
vals = np.take_along_axis(S, idx, axis=1)
order = np.argsort(-vals, axis=1)
np.take_along_axis(idx, order, axis=1) # [[2 0] [1 3]]
np.take_along_axis(vals, order, axis=1) # [[9 5] [8 7]]References
- NumPy docs: User guide (opens in a new tab), API reference (opens in a new tab), Release notes (opens in a new tab) (2.4, 2.5 removals and deprecations)
- NumPy docs: Broadcasting (opens in a new tab), Indexing (opens in a new tab), Copies and views (opens in a new tab): the rules this sheet summarizes
- NumPy docs: Data type promotion (NEP 50) (opens in a new tab), Data types (opens in a new tab)
- NumPy docs: Random
Generator(opens in a new tab),numpy.typing(opens in a new tab),numpy.linalg(opens in a new tab) - NumPy docs: NumPy for MATLAB users (opens in a new tab): side-by-side idioms
- Python Array API standard (opens in a new tab): the portable subset shared with PyTorch, CuPy and JAX
- Nicolas Rougier, From Python to Numpy (opens in a new tab): vectorization techniques
- SciPy (opens in a new tab): sparse matrices, spatial distances, optimization, statistics