Vectorization vs Loops

TL;DR

A Python for loop over an array does the same arithmetic as a NumPy expression. The arithmetic isn’t the bottleneck — the interpreter overhead per element is. Each iteration in pure Python costs ~50–200 nanoseconds in dispatch, attribute lookup, and reference counting. NumPy operations dispatch once to a C / SIMD routine that processes millions of elements without ever touching the interpreter.

The single rule that matters: stay in C. Anything that keeps your data inside NumPy / PyTorch / Numba / a compiled extension is fast. Anything that bounces values out to Python objects per element is slow. The same arithmetic expressed two ways can differ by 100× or more.

# Slow — one Python operation per element
out = []
for x in arr:
    out.append(x * 2 + 3)

# Fast — one NumPy operation, period
out = arr * 2 + 3

This guide is about the patterns that keep you in C — and the fallback (Numba, Cython) for when you genuinely can’t.

The picture in your head

A Python for loop is the interpreter waking up once per iteration: look up the variable, look up __mul__, look up __add__, allocate a new boxed int for the result, append to a list, decrement reference counts, garbage-collect when needed. For a million iterations, that’s a million round-trips through the interpreter.

A NumPy expression like arr * 2 + 3 makes one call into a C routine that knows the dtype, knows the buffer, and processes everything in a tight loop with SIMD instructions. The interpreter only sees: “call ufunc once, get back an ndarray.” Per-element overhead: zero.

The cost difference is the constant factor between “instant” and “go get a coffee.”

A worked benchmark

import numpy as np
import time

n = 1_000_000
arr = np.random.default_rng(0).random(n, dtype=np.float32)

# v1 — pure Python loop
def py_loop(arr):
    out = [0.0] * len(arr)
    for i in range(len(arr)):
        out[i] = arr[i] * 2.0 + 3.0
    return out

# v2 — list comprehension (slightly less overhead per iter)
def list_comp(arr):
    return [x * 2.0 + 3.0 for x in arr]

# v3 — NumPy
def numpy_vec(arr):
    return arr * 2.0 + 3.0

for fn in (py_loop, list_comp, numpy_vec):
    t0 = time.perf_counter()
    fn(arr)
    print(f"{fn.__name__}: {(time.perf_counter() - t0)*1000:.1f} ms")

Indicative timings on a modern laptop:

That ratio is consistent across operations. NumPy is roughly two orders of magnitude faster than a Python loop on million-element arrays. For larger arrays it gets even better, because vectorized code also benefits from SIMD and cache locality.

The vectorization recipe

The translation is usually mechanical. A for loop that:

1. Iterates over an array,
2. Computes one expression per element,
3. Stores the result into another array

becomes a single NumPy expression on whole arrays.

The right-hand-side operations all dispatch to C and are fast.

np.where — the vectorized if

The most common loop-to-vector translation involves a conditional:

# Slow
out = np.empty_like(scores)
for i in range(len(scores)):
    out[i] = 1 if scores[i] > 0.5 else 0

# Fast
out = np.where(scores > 0.5, 1, 0)

# Even faster (no ufunc, just a boolean → int cast)
out = (scores > 0.5).astype(np.int32)

For three-way branches, nest:

out = np.where(scores > 0.8, "high",
        np.where(scores > 0.5, "medium", "low"))

For more than ~3 branches, use np.select:

conds = [scores > 0.8, scores > 0.5, scores > 0.2]
choices = ["high", "medium", "low"]
out = np.select(conds, choices, default="very_low")

Boolean indexing — the vectorized filter

arr[mask] returns the elements where mask is True. Way faster than a loop with an if.

# Slow
positives = []
for x in arr:
    if x > 0:
        positives.append(x)

# Fast
positives = arr[arr > 0]

You can also assign into a boolean slice — the vectorized “in-place update”:

arr[arr < 0] = 0          # ReLU in one line
arr = np.clip(arr, 0, 1)  # also in one line; equivalent for the upper bound too

Aggregation — sum, mean, max, argmax, cumsum

Reductions in NumPy are C loops over the buffer. Use them.

arr.sum()                     # over everything
arr.sum(axis=0)               # column sums for 2D
arr.mean(axis=-1)             # mean over last axis
arr.max()                     # single max
arr.argmax(axis=1)            # index of max along each row
np.cumsum(arr)                # running sum

The pattern that bites: writing a Python loop that maintains running_total = 0; for x in arr: running_total += x; out.append(running_total). Replace with np.cumsum(arr) — same semantics, 100× faster.

When you genuinely can’t vectorize

Some loops resist vectorization:

• State that depends on previous iterations. Each step needs the result of the last. Tree traversals. Iterative simulations.
• Variable-length per-iteration work. “Process each example until convergence” with different convergence times.
• Integration with non-NumPy code. Calling out to a C library that takes one example at a time.

For these, three escape hatches:

np.vectorize — easy but mostly cosmetic

@np.vectorize
def f(x):
    return some_python_function(x)

This is a for-loop wrapper, not a real speedup. It’s there for API convenience (broadcasting, dtype handling), not performance. Don’t reach for it expecting speed.

Numba — JIT compile the Python loop

Numba compiles a Python function to LLVM. If the function is “loop over an array, do scalar arithmetic,” the result is C-fast.

import numba

@numba.njit(cache=True)
def f(arr):
    out = np.empty_like(arr)
    total = 0.0
    for i in range(len(arr)):
        total = total * 0.9 + arr[i]    # iterative state
        out[i] = total
    return out

@numba.njit (no Python mode) gives you C-level speed for code that sticks to a subset of Python (numbers, arrays, basic control flow). First call has a compile cost (~seconds); subsequent calls are microseconds.

Cython / C extensions / Rust via PyO3

For performance-critical inner loops, drop to a compiled language. PyTorch’s torch.compile (@torch.compile) is increasingly the modern path for tensor code: it traces your Python and emits Triton / CUDA kernels.

A worked ML example — softmax

The textbook scalar definition:

def softmax_scalar(x):
    out = []
    total = 0.0
    for v in x:
        total += math.exp(v)
    for v in x:
        out.append(math.exp(v) / total)
    return out

Two passes, slow, and numerically unstable for large inputs.

The NumPy version:

def softmax(x: np.ndarray) -> np.ndarray:
    x = x - x.max(axis=-1, keepdims=True)   # subtract max for stability
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

Three lines, fully vectorized, numerically stable, broadcasts cleanly over any leading batch dimensions. This is the actual implementation shape used inside every neural-network library.

For a (B, N) input, the scalar version is O(B × N) Python operations. The NumPy version is O(B × N) C operations — same big-O, ~100× constant factor. On a (64, 50000) softmax (3.2M values), the NumPy version takes ~5 ms; the loop takes ~600 ms.

Common gotchas

• Vectorizing a sequential operation. If iteration i depends on iteration i-1, you can’t trivially vectorize. Use np.cumsum / np.cumprod for prefix-aggregated forms; otherwise reach for Numba.
• np.vectorize is not a speedup. It’s a syntactic convenience.
• For-loops over array indices. for i in range(len(a)): a[i] = ... is just as slow as for x in a: .... The bottleneck is the Python loop.
• Pandas .apply with a Python function is a Python loop in disguise. Use .transform, .agg, or vectorized expressions instead. See PY 205.
• Mixing dtypes silently upcasts. float32 + float64 → float64, doubling memory and cutting bandwidth. Use .astype deliberately.
• Allocating intermediate arrays. (a - a.mean()) / a.std() builds three temporary arrays. For huge arrays, in-place operations (np.subtract(a, a.mean(), out=a)) save allocation cost.

Where this shows up in real ML code

• Featurisation pipelines. Convert lists/dicts to NumPy, then every transformation is vectorized.
• Loss functions. Implemented in C++/CUDA in PyTorch; the Python wrapper is one line.
• Data augmentation. torchvision.transforms runs vectorized in PIL or tensor mode; never per-pixel Python.
• Simulations / RL. Vectorized environments (gymnasium’s VectorEnv) run 1000s of envs in parallel arrays.
• Tokenization. tokenizers-rs library runs in Rust; the Python wrapper batches inputs to amortise the call overhead.

The defensive habit: when you write for i in range(len(arr)): over a NumPy array, stop and ask “what’s the vectorized version?” 19 times out of 20, there is one, and it’s a one-liner.

Resources

• NumPy docs — Universal functions (ufunc) — numpy.org — what dispatches to C, what doesn’t.
• High Performance Python (2nd ed.) — Gorelick & Ozsvald — oreilly.com — Chapters 6–7 cover vectorization, broadcasting, Numba.
• Numba documentation — numba.readthedocs.io — when you can’t vectorize but need speed.
• From Python to Numpy — Vectorization chapter — labri.fr — pattern catalogue with worked examples.
• PyTorch — torch.compile — pytorch.org — modern equivalent for PyTorch tensor code.