NumPy Guide: Arrays, Broadcasting and Vectorization

Most slow NumPy code is Python code that happens to use arrays: a for loop over rows, one element at a time. Most confusing NumPy errors come from one rule about shapes. This guide covers both. It explains what vectorisation means, states the broadcasting rule precisely, and then works through the three bugs that rule causes in real data code.

What vectorisation means

The NumPy documentation defines vectorisation as "the absence of any explicit looping, indexing, etc., in the code". The loops still happen, but in pre-compiled C inside NumPy instead of in the Python interpreter (NumPy: what is NumPy?). You describe the operation on whole arrays, and NumPy runs the loop.

import numpy as np

prices = np.array([120.0, 80.5, 99.9, 45.0])
qty = np.array([3, 10, 1, 7])

# Loop version: Python executes every iteration
revenue = np.empty(len(prices))
for i in range(len(prices)):
    revenue[i] = prices[i] * qty[i]

# Vectorised version: one expression, the loop runs inside NumPy
revenue = prices * qty          # array([360. , 805. ,  99.9, 315. ])
total = revenue.sum()           # 1579.9

On four elements you will not notice the difference. On millions you will. Measure it yourself with %timeit in Jupyter rather than trusting a quoted speed-up, since the gap depends on the operation and the array size.

The broadcasting rule

Broadcasting is how NumPy combines arrays of different shapes in one operation. There is exactly one rule (NumPy: broadcasting):

1. Write the two shapes right-aligned and compare them from the last dimension, moving left.
2. Two dimensions are compatible when they are equal, or when one of them is 1. A size-1 dimension is stretched to match the other.
3. If one array has fewer dimensions, the missing leading dimensions count as 1.
4. If any pair fails, NumPy raises ValueError: operands could not be broadcast together.

The stretching is conceptual. NumPy does not copy the smaller array to the larger shape; it reuses the same values, which is why broadcasting is memory-efficient in the common cases.

Broadcasting in real data code

Take a feature matrix X with one row per sample and one column per feature, shape (n_samples, n_features). Standardising each column is one line, because the per-column statistics have shape (n_features,) and line up with the last axis:

rng = np.random.default_rng(0)
X = rng.normal(loc=50, scale=10, size=(1000, 3))   # shape (1000, 3)

mu = X.mean(axis=0)                                # shape (3,)
sigma = X.std(axis=0)                              # shape (3,)
Z = (X - mu) / sigma                               # (1000, 3) with (3,) → (1000, 3)

Bug 1: the silent wrong axis

Now normalise each row to sum to 1. The row sums have shape (1000,), which does not line up with the last axis (3), so NumPy raises an error. That is the helpful case. Keep the reduced axis with keepdims=True and the shapes line up:

row_sums = X.sum(axis=1)                  # shape (1000,)
# X / row_sums                            # ValueError: (1000,3) vs (1000,)

row_sums = X.sum(axis=1, keepdims=True)   # shape (1000, 1)
P = X / row_sums                          # (1000, 3) with (1000, 1) → (1000, 3)

The dangerous case is a square array. With a 3 × 3 matrix, M / M.sum(axis=1) does not fail. The (3,) row sums line up with the last axis and each column is divided by a row’s sum. You get a result of the right shape and the wrong values, with no error. Use keepdims=True for every reduction you intend to broadcast back, and the square case is correct too.

M = np.array([[1., 1., 2.],
              [3., 3., 6.],
              [0., 5., 5.]])

wrong = M / M.sum(axis=1)                 # runs, but divides column j by row j's sum
right = M / M.sum(axis=1, keepdims=True)  # each row now sums to 1
right.sum(axis=1)                         # array([1., 1., 1.])

Bug 2: the intermediate that does not fit in memory

np.newaxis (the same as None in an index) inserts a size-1 dimension, which lets you broadcast every row of one array against every row of another. That is exactly what a pairwise distance needs:

A = rng.random((500, 64))    # 500 points, 64 features
B = rng.random((800, 64))    # 800 points

diff = A[:, np.newaxis, :] - B[np.newaxis, :, :]   # (500, 1, 64) - (1, 800, 64) → (500, 800, 64)
dist = np.sqrt((diff ** 2).sum(axis=-1))           # (500, 800)

It is elegant, and it builds a (500, 800, 64) intermediate: 25.6 million float64 values, about 205 MB. Scale both inputs to 10,000 points and the intermediate becomes 6.4 billion values, about 51 GB, before the sum ever runs. The NumPy documentation names this case: broadcasting can use "unnecessarily large amounts of memory", and then it is better to write the outer loop in Python (NumPy: broadcasting). Process A in chunks of rows, or use a routine built for the job, such as SciPy's cdist.

Bug 3: edits that leak through views

Slicing a NumPy array (basic indexing) returns a view: a new array object that shares the original's memory. Indexing with a list or a boolean mask (advanced indexing) always returns a copy (NumPy: copies and views). Writing to a view changes the original:

data = np.arange(10)
first_half = data[:5]        # view
first_half[:] = 0
data                         # array([0, 0, 0, 0, 0, 5, 6, 7, 8, 9]) — original changed

picked = data[[5, 6]]        # advanced indexing → copy
picked[:] = -1
data[5:7]                    # array([5, 6]) — original untouched

safe = data[:5].copy()       # take an explicit copy when you mean to modify

In a pipeline, this shows up as a "cleaning" step that quietly modifies the raw array another step still needs. When a function modifies an array it received, copy it first, or document that it works in place.

When you cannot vectorise

Vectorisation needs each output to be computable independently, or through a standard reduction or accumulation. Ufuncs provide reduce and accumulate, and reductions take an axis argument (NumPy: ufunc basics). A running total that resets whenever a condition holds, or any step that depends on the previous output in a custom way, is inherently sequential. Keep that loop, make it as small as possible, and vectorise everything around it.

Where to go next

Apply the same thinking to tables in the pandas guide to DataFrames, cleaning and aggregation. Test yourself with NumPy interview questions for data roles. NumPy and pandas are covered in the data-handling module of our Python programming course.

Sources

• Broadcasting — NumPy user guide
• What is NumPy? — NumPy user guide
• Copies and views — NumPy user guide
• numpy.vectorize — NumPy API reference
• Universal functions (ufunc) basics — NumPy user guide