NumPy: The Numerical Engine

Python's built-in lists are pointer arrays of arbitrary objects — flexible, but ruinously slow for numerical work. NumPy throws that away. An ndarray is a single block of contiguous memory, every element the same fixed-width type, with a small header describing shape and stride. That tiny shift in representation unlocks vectorised C loops, SIMD instructions, and BLAS-backed linear algebra. The difference is often 50–200× over equivalent Python.

You've seen the basics in ds-numpy. This lesson goes deep — memory layout, broadcasting rules, advanced indexing, axis semantics, and the performance discipline that separates working NumPy from production NumPy. By the end you should be writing array expressions reflexively, never reaching for a for loop over rows.

1. The Memory Model — Why NumPy Is Fast

An ndarray is three things bundled together:

1. A data buffer — a flat C array of identical-width values.
2. A dtype — the element type (int32, float64, bool, etc.) and its byte width.
3. A shape + strides tuple — how to interpret the buffer as an n-dimensional grid.

python

import numpy as np

a = np.arange(12, dtype=np.int32).reshape(3, 4)
print(a)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

print(a.dtype)         # int32
print(a.shape)         # (3, 4)
print(a.strides)       # (16, 4)   — 16 bytes to next row, 4 bytes to next column
print(a.nbytes)        # 48        — 12 ints * 4 bytes each

strides=(16, 4) is the key. Moving down a row jumps 16 bytes (one full row of four int32s); moving across a column jumps 4 bytes (one int32). Reshaping, slicing, and transposing usually just change the strides — no data is copied. That's why a.T is instantaneous on a 1 GB array.

Compare with a Python list of lists: each row is a separate PyListObject, each cell is a separate PyLongObject somewhere else on the heap, every access dereferences three pointers. NumPy's memory locality is what makes the CPU happy.

2. Creating Arrays Efficiently

Avoid np.array([0]*1000) — that builds a Python list first. The right constructors allocate the buffer directly:

python

import numpy as np

np.zeros((3, 4), dtype=np.float32)      # filled with 0.0, 12 floats = 48 bytes
np.ones((2, 2), dtype=np.int64)         # filled with 1
np.full((2, 3), 7.5)                    # filled with 7.5
np.empty((3, 3))                        # ALLOCATES — values are garbage; only use if you'll overwrite
np.eye(4)                               # 4x4 identity matrix
np.arange(0, 10, 0.5)                   # like range, but with float step — returns array
np.linspace(0, 1, 11)                   # 11 evenly spaced points from 0 to 1 inclusive

dtype choice matters for memory. A (1000, 1000) float64 array is 8 MB; the same shape as float32 is 4 MB; as int8 it's 1 MB. Pick the narrowest type that holds your data — but only with intent. Silent integer overflow is real (see Common Mistakes).

For modern random data, use the Generator API — not the legacy np.random.rand / randn global state:

python

rng = np.random.default_rng(seed=42)     # the modern API (NEP 19)
rng.standard_normal((3, 4))              # 3x4 of N(0, 1)
rng.uniform(0, 100, size=1000)           # 1000 uniform on [0, 100)
rng.integers(0, 10, size=20)             # 20 ints in [0, 10)
rng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])

+ setup added so this can run · defines np

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

The legacy functions still work but share a global seed that's hostile to parallelism. Use default_rng in new code.

3. Vectorisation — Don't Loop, Express

The single rule for writing fast NumPy: state your computation as an array operation, never as a per-element Python loop. Element-wise arithmetic, comparison, and unary functions are all vectorised:

python

import numpy as np

a = np.array([1, 2, 3, 4])
b = np.array([10, 20, 30, 40])

a + b                       # [11, 22, 33, 44]   — element-wise
a * b                       # [10, 40, 90, 160]
a ** 2                      # [1, 4, 9, 16]
np.sqrt(a)                  # [1., 1.41, 1.73, 2.]
np.exp(a)                   # [2.72, 7.39, 20.09, 54.60]
a > 2                       # [False, False, True, True]   — boolean array

# The "I want a Python loop" trap
result = np.empty_like(a)
for i in range(len(a)):
    result[i] = a[i] ** 2 + b[i]      # SLOW — Python loop over a C array

# The vectorised way
result = a ** 2 + b                   # FAST — one C-level pass

A vectorised expression on a million-element array runs in milliseconds; the equivalent Python loop takes seconds. Every operator and every function in np.* is built to work on whole arrays at once.

4. Broadcasting in Depth

Broadcasting is the rule that lets arrays of different shapes combine without explicit replication. NumPy aligns shapes from the right, then for each axis demands one of three things: the dimensions match, one is 1 (stretched), or one is missing (prepended as 1).

The rules

Operation	Result shape	Reasoning
`(3, 4)` + `(4,)`	`(3, 4)`	RHS becomes `(1, 4)`, stretched to `(3, 4)`
`(3, 1)` * `(1, 4)`	`(3, 4)`	Both stretched to `(3, 4)` — outer product
`(3, 4)` + `(3,)`	error	Right-align: `4 != 3` and neither is 1
`(3, 4)` + `(3, 1)`	`(3, 4)`	Column vector stretched across columns
`(5,)` + `5.0`	`(5,)`	Scalar broadcasts to anything

No memory is allocated for the "stretching" — NumPy iterates with stride zero on the broadcast axis. Free virtual replication.

Common patterns

python

import numpy as np

# Subtract the column mean from each column (centre per feature)
X = np.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0],
              [7.0, 8.0, 9.0]])

col_mean = X.mean(axis=0)               # shape (3,)
centred  = X - col_mean                 # (3,3) - (3,) -> (3,3), each column zero-mean

# Z-score normalisation per column
std = X.std(axis=0)
z = (X - col_mean) / std

# Outer product via broadcasting (no np.outer needed)
u = np.array([1, 2, 3])                 # shape (3,)
v = np.array([10, 20, 30, 40])          # shape (4,)
outer = u[:, None] * v[None, :]         # (3,1) * (1,4) -> (3,4)
# [[ 10  20  30  40]
#  [ 20  40  60  80]
#  [ 30  60  90 120]]

# Pairwise distances (Euclidean) between rows of A and rows of B
A = np.random.default_rng(0).standard_normal((100, 5))
B = np.random.default_rng(1).standard_normal((50, 5))
# (100,1,5) - (1,50,5) -> (100,50,5), sum-square last axis -> (100,50)
dists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))
print(dists.shape)                      # (100, 50)

X[:, None] and X[None, :] are the workhorse idioms — they add a new axis of length 1 so the broadcasting machinery can take over.

5. Advanced Indexing

Beyond a[0] and a[1:5] lie three more flexible indexing modes.

Integer-array indexing — pick arbitrary rows or columns

python

import numpy as np

a = np.arange(20).reshape(5, 4)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]
#  [12 13 14 15]
#  [16 17 18 19]]

a[[0, 2, 4]]                # rows 0, 2, 4 — shape (3, 4)
a[:, [3, 1, 0]]             # columns reordered — shape (5, 3)
a[[0, 1, 2], [3, 2, 1]]     # element-wise: (0,3), (1,2), (2,1) -> [3, 6, 9]

Boolean masks — filter by condition

python

mask = a > 10
a[mask]                     # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19]
a[a % 2 == 0]               # all evens

# Mask assignment — set every negative entry to zero
x = np.array([-3, 1, -1, 5, -2])
x[x < 0] = 0
print(x)                    # [0 1 0 5 0]

+ setup added so this can run · defines a, np

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

a = _AutoMock('a')
np = _AutoMock('np')

`np.where` — vectorised ternary

python

x = np.array([-3, 1, -1, 5, -2])
np.where(x < 0, 0, x)       # [0, 1, 0, 5, 0]   — same as max(x, 0) here
np.where(x < 0, "neg", "pos")
# array(['neg', 'pos', 'neg', 'pos', 'neg'], dtype='<U3')

# Two-argument form returns indices where condition is true
idx = np.where(x < 0)
print(idx)                  # (array([0, 2, 4]),)
print(x[idx])               # [-3 -1 -2]

+ setup added so this can run · defines np

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

Fancy indexing always creates a copy, never a view (more on that below). Boolean masks also copy. Plain slices (a[1:5]) are views — modifying them mutates the original.

6. Aggregations and the `axis` Argument

Reductions like sum, mean, std, min, max, argmax all accept axis=. The intuition: axis=k collapses axis k, leaving the others.

python

import numpy as np

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

a.sum()                       # 21    — collapses everything
a.sum(axis=0)                 # [5, 7, 9]   — collapse rows, keep columns -> per-column sum
a.sum(axis=1)                 # [6, 15]     — collapse columns, keep rows -> per-row sum

a.mean(axis=0)                # [2.5, 3.5, 4.5]   — column means
a.argmax(axis=1)              # [2, 2]            — index of max in each row

Mnemonic — axis=0 is "down through the rows" (gives per-column results); axis=1 is "across the columns" (gives per-row results). It feels backwards until it doesn't.

keepdims=True preserves the collapsed axis as length 1 — essential for broadcasting back:

python

row_max = a.max(axis=1)                  # shape (2,)   — can't subtract from (2,3)
row_max = a.max(axis=1, keepdims=True)   # shape (2, 1) — broadcasts cleanly
normalised = a - row_max                 # subtract per-row max — log-softmax trick

+ setup added so this can run · defines a

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

a = _AutoMock('a')

7. Linear Algebra — `@`, `dot`, `solve`

@ is the matrix-multiplication operator (PEP 465). For higher-dimensional arrays it broadcasts in batched-matmul fashion.

python

import numpy as np

A = np.array([[1, 2], [3, 4]], dtype=float)
b = np.array([5, 6], dtype=float)

A @ b                           # [17., 39.]    — matrix-vector
A @ A                           # 2x2 matrix-matrix
A.T @ A                         # Gram matrix

# Solve Ax = b — DO NOT compute inv(A) @ b
x = np.linalg.solve(A, b)       # [-4., 4.5]    — faster, more stable

# When you really need the inverse
A_inv = np.linalg.inv(A)

# Norms — Euclidean, Manhattan, Frobenius
np.linalg.norm(b)               # 7.81   — L2 by default
np.linalg.norm(b, ord=1)        # 11.0   — L1
np.linalg.norm(A, ord="fro")    # 5.47   — Frobenius

# Eigendecomposition
eigvals, eigvecs = np.linalg.eig(A)
# eigvals: [-0.37,  5.37]
# eigvecs: columns are the corresponding unit eigenvectors

# Singular value decomposition (PCA's engine)
U, s, Vt = np.linalg.svd(A, full_matrices=False)

For Ax = b, always use solve over inv @ b — it uses LU decomposition under the hood, is ~2× faster, and avoids the numerical instability of explicit inversion.

8. Statistics and Random — Beyond the Basics

python

import numpy as np

data = np.array([2, 4, 4, 4, 5, 5, 7, 9])

np.percentile(data, [25, 50, 75])       # [4., 4.5, 5.5]   — quartiles
np.median(data)                         # 4.5
np.cov(data)                            # 4.57   — sample covariance (1-D = variance)

# Correlation matrix
x = np.random.default_rng(0).standard_normal((100,))
y = 2 * x + 0.5 * np.random.default_rng(1).standard_normal((100,))
np.corrcoef(x, y)
# [[1.   0.97]
#  [0.97 1.  ]]

# Multi-dimensional stats — pick the axis carefully
X = np.random.default_rng(0).standard_normal((1000, 5))   # 1000 samples, 5 features
feature_mean = X.mean(axis=0)           # shape (5,)   — mean per feature
feature_std  = X.std(axis=0, ddof=1)    # sample std (N-1 denominator)

ddof=1 matters for unbiased sample statistics — NumPy defaults to population (ddof=0). Pandas defaults to sample. The difference burns people.

9. Reshape, Transpose, Stack, Split

Shape manipulation almost never copies data — it just reinterprets the same buffer.

python

import numpy as np

a = np.arange(12)
a.reshape(3, 4)                 # 3x4 view
a.reshape(2, -1)                # -1 means "infer" — 2x6
a.reshape(2, 2, 3)              # 3-D view

# Transpose — swap or reorder axes
M = np.arange(24).reshape(2, 3, 4)
M.T.shape                       # (4, 3, 2)   — reverse all axes
M.transpose(1, 0, 2).shape      # (3, 2, 4)   — explicit reorder

# Stacking and concatenation
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

np.concatenate([a, b])          # [1, 2, 3, 4, 5, 6]
np.stack([a, b])                # [[1,2,3],[4,5,6]]   — new axis at front
np.vstack([a, b])               # [[1,2,3],[4,5,6]]   — row-stack
np.hstack([a, b])               # [1, 2, 3, 4, 5, 6]  — column-stack (1-D)
np.column_stack([a, b])         # [[1,4],[2,5],[3,6]] — pair as columns

# Splitting — inverse of stack
M = np.arange(12).reshape(3, 4)
np.split(M, 2, axis=1)          # split into 2 along columns -> list of (3,2) arrays

reshape requires the new shape to have the same total size. ravel() flattens to 1-D (view if contiguous); flatten() always copies.

10. Views vs Copies — The Trap

Slicing returns a view — same buffer, different shape/stride. Modifying it mutates the original. Fancy indexing returns a copy — independent data.

python

import numpy as np

a = np.arange(10)
b = a[2:6]                      # VIEW — shares memory with a
b[0] = 999
print(a)                        # [  0   1 999   3   4   5   6   7   8   9]   — surprise!

c = a[[2, 3, 4]]                # COPY — fancy indexing
c[0] = -1
print(a)                        # unchanged

# Check explicitly
print(b.base is a)              # True   — b is a view of a
print(c.base is a)              # False  — c is independent

# Force a copy when you want one
d = a[2:6].copy()
d[0] = 555
print(a)                        # unchanged

This is the cause of countless "why did my data change?" bugs. When you intend to mutate, .copy() defensively. When you intend a view (for memory or speed), confirm with .base.

11. Performance — `np.einsum`, Benchmarks, and the Loop Tax

The benchmark that ends every "should I vectorise?" debate:

python

import numpy as np, time

n = 1_000_000
a = np.random.default_rng(0).standard_normal(n)
b = np.random.default_rng(1).standard_normal(n)

# Python loop
t0 = time.perf_counter()
result_loop = [a[i] * b[i] + a[i] for i in range(n)]
t_loop = time.perf_counter() - t0

# NumPy vectorised
t0 = time.perf_counter()
result_vec = a * b + a
t_vec = time.perf_counter() - t0

print(f"loop: {t_loop:.3f}s   vec: {t_vec:.4f}s   speedup: {t_loop/t_vec:.0f}x")
# loop: ~0.300s   vec: ~0.003s   speedup: ~100x

For complex multi-axis operations, np.einsum expresses tensor contractions in Einstein notation — sometimes faster than equivalent combinations of transpose, reshape, and @:

python

import numpy as np

A = np.random.default_rng(0).standard_normal((100, 50))
B = np.random.default_rng(1).standard_normal((50, 30))

# Standard matmul
np.einsum("ij,jk->ik", A, B).shape          # (100, 30)   — same as A @ B

# Batched dot product of corresponding rows
X = np.random.default_rng(0).standard_normal((1000, 5))
Y = np.random.default_rng(1).standard_normal((1000, 5))
dots = np.einsum("ij,ij->i", X, Y)          # shape (1000,) — row-wise dot

When the GIL or NumPy's vectorised primitives aren't enough — a hot inner loop with branching that can't be expressed in array ops — reach for numba (@njit JIT-compiles a Python function) or Cython. Both bypass the Python interpreter for the hot path.

12. I/O — Saving and Loading

python

import numpy as np

a = np.arange(12).reshape(3, 4)

np.save("a.npy", a)                         # binary, fastest, dtype-preserving
b = np.load("a.npy")

np.savez("bundle.npz", x=a, y=a * 2)        # multiple arrays in one file
data = np.load("bundle.npz")
print(data["x"], data["y"])

np.savetxt("a.csv", a, delimiter=",", fmt="%d")   # human-readable, lossy for floats
b = np.loadtxt("a.csv", delimiter=",")

.npy is the default for "save this for later" — round-trips exactly, including dtype, shape, and endianness. .csv is for handoff to non-NumPy consumers.

Common Mistakes

1. Silent integer overflow

python

import numpy as np

a = np.array([100, 100], dtype=np.int8)     # range [-128, 127] — 100 fits
print(a + a)                                # [-56 -56]  — WRAPPED, no warning

NumPy integer dtypes wrap silently on overflow. int8, int16, int32 all bite this way. Stick with int64 (the default on 64-bit platforms) unless memory pressure forces you down. If you go smaller, validate ranges at the boundary.

2. Mixed dtypes producing object arrays

python

a = np.array([1, 2, "three", 4])
print(a.dtype)                              # <U21   — Unicode strings, slow

a = np.array([1, 2, None, 4])
print(a.dtype)                              # object — boxed Python objects, slowest

+ setup added so this can run · defines np

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

The moment NumPy can't find a common numeric dtype, it falls back to object (a pointer array). You lose every vectorisation benefit — it's slower than a list of lists. Clean missing data with np.nan (which is a float) before constructing the array, or use pandas for genuinely heterogeneous data.

3. In-place ops on shared views

python

import numpy as np

a = np.arange(10)
b = a[::2]                  # view of every other element: [0, 2, 4, 6, 8]
b *= 100                    # IN-PLACE — also mutates a
print(a)                    # [  0   1 200   3 400   5 600   7 800   9]

*=, +=, etc. modify the underlying buffer. If b is a view, a sees the change. Use b = b * 100 (creates a new array) if you didn't mean to propagate, or .copy() before the in-place op.

4. Reaching for Python operators when NumPy ops exist

python

import math
import numpy as np

x = np.linspace(0, 10, 1_000_000)

# Slow — Python loop hidden inside a list comprehension
y = [math.sin(v) for v in x]                # ~200ms

# Fast — vectorised C loop
y = np.sin(x)                               # ~5ms

math.sin, math.exp, math.log are scalar functions. np.sin, np.exp, np.log are vectorised. Mixing them undoes the point of using NumPy. The same applies to min/max — np.min(a) is much faster than min(a) on a large array.

5. Forgetting keepdims and breaking broadcast

python

import numpy as np

a = np.random.default_rng(0).standard_normal((100, 5))

# Wrong — shape mismatch
centred = a - a.mean(axis=1)                # (100, 5) - (100,) — fails to align

# Right — preserve the collapsed axis as length 1
centred = a - a.mean(axis=1, keepdims=True) # (100, 5) - (100, 1) — broadcasts cleanly

Whenever you reduce and then combine, keepdims=True is almost always what you want.

🎯 Your Turn — Rolling Mean in O(n)

Write moving_average(arr, window) that returns the rolling mean of arr with a fixed window size. The naive solution is O(n × window); the right one is O(n) using np.cumsum.

Expected behaviour:

python

moving_average(np.array([1, 2, 3, 4, 5, 6, 7]), window=3)
# array([2., 3., 4., 5., 6.])           — means of [1,2,3], [2,3,4], ..., [5,6,7]

moving_average(np.array([1.0, 2.0, 4.0, 8.0]), window=2)
# array([1.5, 3.0, 6.0])

+ setup added so this can run · defines moving_average, np

# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

def moving_average(*_a, **_kw):
    print('-> moving_average() called')
    return _AutoMock('moving_average()')
np = _AutoMock('np')

Constraints:

Return a 1-D ndarray of length len(arr) - window + 1.
Must run in O(n) — single pass — using np.cumsum, not a Python loop.
Raise ValueError if window < 1 or window > len(arr).

Skeleton:

python

import numpy as np

def moving_average(arr, window):
    # TODO 1: validate window (>= 1, <= len(arr))
    # TODO 2: compute cumulative sum, prepended with a 0
    # TODO 3: subtract cumsum[:-window] from cumsum[window:] to get window sums
    # TODO 4: divide by window to get means
    ...

Hint 1 — Cumsum gives you prefix sums

np.cumsum(arr) returns [arr[0], arr[0]+arr[1], arr[0]+arr[1]+arr[2], ...]. The sum of any window arr[i:i+w] equals cumsum[i+w-1] - cumsum[i-1]. Prepend a zero (np.concatenate([[0], cumsum])) so the i=0 case works without a special branch.

Hint 2 — Vectorise the subtraction

With the prepended cumsum (length n+1), the array of window sums is cumsum[window:] - cumsum[:-window]. Both are length n - window + 1. Divide by window (a scalar, broadcasts automatically) to get means. No loop.

Show full solution

python

import numpy as np

def moving_average(arr, window):
    """Rolling mean of `arr` over a fixed-size window. O(n) via cumsum."""
    if window < 1:
        raise ValueError("window must be >= 1")
    if window > len(arr):
        raise ValueError("window cannot exceed array length")

    cumsum = np.cumsum(arr, dtype=np.float64)            # prefix sums
    cumsum = np.concatenate(([0.0], cumsum))             # prepend 0 — so cumsum[0] == 0
    window_sums = cumsum[window:] - cumsum[:-window]     # sums of every length-w window
    return window_sums / window


# Verify
print(moving_average(np.array([1, 2, 3, 4, 5, 6, 7]), 3))
# [2. 3. 4. 5. 6.]

print(moving_average(np.array([1.0, 2.0, 4.0, 8.0]), 2))
# [1.5 3.  6. ]

# Performance — naive vs cumsum on a million elements
import time
arr = np.random.default_rng(0).standard_normal(1_000_000)
w = 100

t0 = time.perf_counter()
naive = np.array([arr[i:i+w].mean() for i in range(len(arr) - w + 1)])
t_naive = time.perf_counter() - t0

t0 = time.perf_counter()
fast = moving_average(arr, w)
t_fast = time.perf_counter() - t0

print(f"naive: {t_naive:.3f}s   cumsum: {t_fast:.4f}s   speedup: {t_naive/t_fast:.0f}x")
# naive: ~1.2s   cumsum: ~0.01s   speedup: ~100x

Why this is O(n) — cumsum is a single linear pass. The subtraction and division are also linear, vectorised in C. No nested loop hides anywhere. Doubling the window size doesn't change the runtime.

Why dtype=np.float64 on cumsum — if arr is an integer dtype, cumulative sums can overflow silently. Promoting to float64 up front avoids that. For very long arrays of large floats, even float64 accumulates error — np.cumsum uses naive summation. Production-grade rolling means use Kahan summation or the algorithm in pandas' rolling().mean(), which subtracts the leaving element rather than recomputing from a global prefix sum (better numerics, same O(n)).

This trick — prefix sums for any sliding-window aggregate — generalises far. Rolling sum, rolling mean, "is there a window of size k summing to at least X?", "longest subarray with sum < S" — all reduce to operations on a cumsum array. Worth burning into muscle memory.

What You Learned

An ndarray is a contiguous buffer + dtype + shape + strides — reshaping and slicing rearrange metadata, not data.
Use the right constructor (zeros, empty, full, eye, linspace) and the modern default_rng for random data.
Vectorise — express computations as whole-array operations, never per-element Python loops. Expect 50–200× speedups.
Broadcasting aligns shapes from the right; size-1 and missing dimensions stretch. Use [:, None] and keepdims=True to control alignment.
Advanced indexing — integer arrays and boolean masks both copy; plain slices are views. np.where for vectorised ternaries.
axis=k collapses axis k. axis=0 gives per-column results, axis=1 per-row.
Linear algebra: prefer @ over np.dot, np.linalg.solve(A, b) over inv(A) @ b.
Views vs copies — slices are views, fancy indexing copies. Mutations on views propagate. Use .copy() defensively.
Watch for silent integer overflow, accidental object dtype, and the in-place-on-shared-view trap.
For loops the GIL can't help, reach for numba or Cython; for tensor contractions, np.einsum.

Next: Pandas DataFrames: Command Center — the same array engine wrapped in a labelled, heterogeneous, group-aware container.

1 · The lesson

1. The Memory Model — Why NumPy Is Fast

2. Creating Arrays Efficiently

3. Vectorisation — Don't Loop, Express

4. Broadcasting in Depth

The rules

Common patterns

5. Advanced Indexing

Integer-array indexing — pick arbitrary rows or columns

Boolean masks — filter by condition

`np.where` — vectorised ternary

6. Aggregations and the `axis` Argument

7. Linear Algebra — `@`, `dot`, `solve`

8. Statistics and Random — Beyond the Basics

9. Reshape, Transpose, Stack, Split

10. Views vs Copies — The Trap

11. Performance — `np.einsum`, Benchmarks, and the Loop Tax

12. I/O — Saving and Loading

Common Mistakes

🎯 Your Turn — Rolling Mean in O(n)

What You Learned

Practice this

NumPy: The Numerical Engine

1. The Memory Model — Why NumPy Is Fast

2. Creating Arrays Efficiently

3. Vectorisation — Don't Loop, Express

4. Broadcasting in Depth

The rules

Common patterns

5. Advanced Indexing

Integer-array indexing — pick arbitrary rows or columns

Boolean masks — filter by condition

np.where — vectorised ternary

6. Aggregations and the axis Argument

7. Linear Algebra — @, dot, solve

8. Statistics and Random — Beyond the Basics

9. Reshape, Transpose, Stack, Split

10. Views vs Copies — The Trap

11. Performance — np.einsum, Benchmarks, and the Loop Tax

12. I/O — Saving and Loading

Common Mistakes

🎯 Your Turn — Rolling Mean in O(n)

What You Learned

`np.where` — vectorised ternary

6. Aggregations and the `axis` Argument

7. Linear Algebra — `@`, `dot`, `solve`

11. Performance — `np.einsum`, Benchmarks, and the Loop Tax