3. Maths & broadcasting

Intermediate · 10 min read

3.1 Vectorised maths — no loops

Operators work on every element at once.

import numpy as np

a = np.array([1, 2, 3, 4])
print(a * 10)
print(a + a)
print(a ** 2)
print(np.sqrt(a).round(3))
print(a / a.sum())          # each value's share of the total

Output

[10 20 30 40]
[2 4 6 8]
[ 1  4  9 16]
[1.    1.414 1.732 2.   ]
[0.1 0.2 0.3 0.4]

Compare with a loop — same result, but slower and longer:

squares_loop = [x ** 2 for x in a]          # Python loop
squares_vec = a ** 2                        # NumPy: one fast operation
print(list(squares_vec) == squares_loop)

Output

True

3.2 Aggregations and axis

m = np.array([[1, 2, 3],
              [4, 5, 6]])
print(m.sum())            # everything
print(m.sum(axis=0))      # down the rows → one value per column
print(m.sum(axis=1))      # across the columns → one value per row
print(m.mean(axis=1))
print(m.max(), m.min(), m.std().round(3))

Output

21
[5 7 9]
[ 6 15]
[2. 5.]
6 1 1.708

Remembering axis

axis=0 collapses the rows (result has one value per column); axis=1 collapses the columns (one value per row). For a batch of embeddings (n_docs, dims), axis=1 works per document.

3.3 Broadcasting — different shapes, one operation

NumPy stretches the smaller array to match the bigger one when the shapes are compatible.

m = np.arange(6).reshape(2, 3)
print(m + 10)                         # scalar → added to every value
print(m + np.array([100, 200, 300]))  # row (3,) → added to each row
print(m * np.array([[1], [-1]]))      # column (2, 1) → multiplies each row

Output

[[10 11 12]
 [13 14 15]]
[[100 201 302]
 [103 204 305]]
[[ 0  1  2]
 [-3 -4 -5]]

The rule: compare shapes from the right. Two sizes fit if they are equal or one of them is 1.

Shapes	Works?	Why
(2, 3) + (3,)	✅	3 = 3; the row is repeated for both rows
(2, 3) + (2, 1)	✅	1 stretches to 3
(2, 3) + (2,)	❌	3 vs 2 — neither is 1

try:
    m + np.array([1, 2])
except ValueError as err:
    print("ValueError:", err)

Output

ValueError: operands could not be broadcast together with shapes (2,3) (2,) 

A practical use — centre each column (subtract the column mean):

data = np.array([[1.0, 200.0], [3.0, 400.0], [5.0, 600.0]])
centred = data - data.mean(axis=0)          # (3, 2) - (2,)
print(centred)

Output

[[  -2. -200.]
 [   0.    0.]
 [   2.  200.]]

3.4 Sorting, ranking and finding the best

scores = np.array([0.42, 0.91, 0.15, 0.77])
print(np.sort(scores))               # sorted values
print(np.argsort(scores))            # indices that would sort it
print(np.argsort(-scores)[:2])       # indices of the 2 highest
print(scores.argmax(), scores.max())
print(np.unique(np.array([3, 1, 3, 2, 1]), return_counts=True))

Output

[0.15 0.42 0.77 0.91]
[2 0 3 1]
[1 3]
1 0.91
(array([1, 2, 3]), array([2, 1, 2]))

Practice

• Create a 4 × 3 array and divide each row by its own sum (use axis and broadcasting).
• Find the index of the smallest value in every row of a 2-D array.

Next: Linear algebra →