Vectorisation & broadcasting

Vectorisation means expressing a calculation over whole arrays instead of looping element by element. NumPy then runs the loop in compiled C, typically 20–100x faster than Python, with shorter and clearer code.

From loop to expression

prices = np.array([100, 250, 400])
qty = np.array([2, 1, 3])
revenue = prices * qty            # element-wise
total = revenue.sum()

# Broadcasting: (3,1) with (1,3) -> (3,3)
col = np.array([[1], [2], [3]])
row = np.array([10, 20, 30])
table = col * row

The loop version needs four lines and a counter; the vectorised version reads like the formula on the whiteboard.

The broadcasting rules

When shapes differ, NumPy lines them up from the right and, for each pair of dimensions, requires that they are equal or that one of them is 1. A dimension of size 1 is stretched (without copying memory) to match.

Shape A Shape B Result Why
(3, 4) (4,) (3, 4) B is treated as (1, 4) and repeated for each row
(3, 4) (3, 1) (3, 4) B's single column is repeated across columns
(3, 1) (1, 4) (3, 4) Both stretch — the outer-product pattern above
(3, 4) (3,) Error 4 and 3 do not match and neither is 1

The fix for the last row is b[:, np.newaxis], which turns (3,) into (3, 1).

A realistic example

Monthly sales for three stores, and a monthly target for each store:

sales = np.array([[12, 15, 11, 20],     # ₹ lakh, 3 stores x 4 months
                  [ 9, 10, 14, 13],
                  [18, 17, 16, 25]])
target = np.array([13, 12, 20])          # one target per store, shape (3,)

pct = sales / target[:, np.newaxis] * 100    # (3, 4) / (3, 1)
print(pct.round(0))
print((pct >= 100).sum(axis=1))              # months each store met target

Without [:, np.newaxis] this raises ValueError: operands could not be broadcast together.

Useful vectorised functions

np.mean, np.std, np.percentile, np.where, np.unique, np.dot, np.linalg.inv, plus np.cumsum, np.diff, np.clip and np.round.

gst_rate = np.array([0.05, 0.12, 0.18])
line_total = prices * qty
payable = np.round(line_total * (1 + gst_rate), 2)
growth = np.diff(sales, axis=1) / sales[:, :-1] * 100    # month-on-month %

Why it is faster

A Python loop does three slow things per element: a bytecode dispatch, a type check and an object allocation. A vectorised call does the type check once, then runs a tight machine-code loop over packed memory, often using SIMD instructions. This is why arr.sum() beats sum(arr) by a wide margin.

Watch the memory

Broadcasting does not copy, but intermediate results do allocate. (a + b) * c on three 500 MB arrays creates a 500 MB temporary. For large data use in-place operators (a += b, np.multiply(a, b, out=a)) or process in chunks.

Common mistakes

  • Looping with for i in range(len(a)) when a single expression would do.
  • Assuming any two shapes broadcast; check by printing a.shape, b.shape.
  • Forgetting keepdims=True or np.newaxis when dividing rows by their totals.
  • Believing an operation broadcast correctly because it ran — a (3,) against (3, 3) gives a valid but wrong answer if you meant columns.

Practice

Build a 5 x 12 array of monthly sales for five branches and a vector of five annual targets. Compute each branch's percentage of target per month with broadcasting, then subtract every month's cross-branch mean from the table to see who is above or below the network average.

B

Written by Bablu Kumar Jha

Founder, BeeProton AI · Mumbai, Maharashtra, India