Broadcasting for distance matrices in NumPy - Step-by-Step Execution

Concept Flow - Broadcasting for distance matrices

Start with two arrays: A and B

↓

Reshape arrays to add extra dimensions

↓

Broadcast arrays to compatible shapes

↓

Compute element-wise differences

↓

Square differences and sum over last axis

↓

Result: Distance matrix between points in A and B

We reshape arrays to add dimensions, then numpy broadcasts them to compute all pairwise distances without loops.

Execution Sample

NumPy

import numpy as np
A = np.array([[1,2],[3,4]])
B = np.array([[5,6],[7,8]])
dist = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=2))

Calculate Euclidean distances between each point in A and each point in B using broadcasting.

Execution Table

Step	Operation	Shape of A	Shape of B	Broadcasted Shape	Result Description
1	Original A	(2, 2)	-	-	A has 2 points with 2 coordinates each
2	Original B	-	(2, 2)	-	B has 2 points with 2 coordinates each
3	Reshape A to A[:, None, :]	(2, 1, 2)	-	-	Add new axis to A for broadcasting
4	Reshape B to B[None, :, :]	-	(1, 2, 2)	-	Add new axis to B for broadcasting
5	Broadcast A and B	(2, 1, 2)	(1, 2, 2)	(2, 2, 2)	Arrays broadcast to shape (2,2,2) for pairwise subtraction
6	Compute difference A - B	(2, 2, 2)	(2, 2, 2)	(2, 2, 2)	Element-wise difference between points
7	Square differences	(2, 2, 2)	-	(2, 2, 2)	Square each coordinate difference
8	Sum over last axis	(2, 2, 2)	-	(2, 2)	Sum squared differences to get squared distances
9	Square root	(2, 2)	-	(2, 2)	Take sqrt to get Euclidean distances
10	Result	-	-	(2, 2)	Distance matrix: dist[i,j] is distance between A[i] and B[j]
11	Exit	-	-	-	All pairwise distances computed, no loops needed

💡 Broadcasting completes pairwise distance calculation without explicit loops.

Variable Tracker

Variable	Start	After Reshape	After Broadcast	After Difference	After Square	After Sum	Final
A	shape (2,2)	shape (2,1,2)	broadcasted to (2,2,2)	values broadcasted	values squared	values summed	unchanged shape (2,1,2)
B	shape (2,2)	shape (1,2,2)	broadcasted to (2,2,2)	values broadcasted	values squared	values summed	unchanged shape (1,2,2)
dist	undefined	undefined	undefined	undefined	undefined	shape (2,2)	final distance matrix

Key Moments - 3 Insights

Why do we add new axes to A and B before subtracting?

Adding new axes changes shapes so numpy can broadcast arrays to a common shape (2,2,2), allowing pairwise subtraction without loops (see execution_table rows 3-5).

How does numpy know to subtract every point in A from every point in B?

Because after reshaping, A and B have shapes that broadcast to (2,2,2), so subtraction happens element-wise across all pairs (execution_table row 6).

Why do we sum over axis=2 after squaring differences?

Axis 2 corresponds to coordinates; summing squares over this axis gives squared Euclidean distance between points (execution_table row 8).

Visual Quiz - 3 Questions

Look at the execution table at step 5. What is the broadcasted shape of A and B?

A (2, 2, 2)

B (2, 1, 2)

C (1, 2, 2)

D (2, 2)

💡 Hint

Check the 'Broadcasted Shape' column in execution_table row 5.

At which step do we sum squared differences over coordinates?

A Step 7

B Step 8

C Step 6

D Step 9

💡 Hint

Look for 'Sum over last axis' in the 'Operation' column of execution_table.

If we skip adding new axes to A and B, what happens to broadcasting?

A Broadcasting still works correctly

B Distance matrix is computed but with wrong shape

C Shapes mismatch, subtraction fails

D Result is a scalar

💡 Hint

Refer to execution_table rows 3 and 4 where reshaping enables broadcasting.

Concept Snapshot

Broadcasting for distance matrices:
- Reshape arrays to add axes (e.g., A[:, None, :])
- Numpy broadcasts shapes to align dimensions
- Compute element-wise differences for all pairs
- Square, sum over coordinates, then sqrt
- Result is a matrix of distances without loops

Full Transcript

This example shows how to compute a distance matrix between two sets of points using numpy broadcasting. We start with two arrays A and B, each holding points as rows. We reshape A and B to add new axes, enabling numpy to broadcast them to a common shape. This lets us subtract every point in B from every point in A at once. Then we square the differences, sum over the coordinate axis to get squared distances, and take the square root to get Euclidean distances. The final result is a matrix where each element is the distance between a point in A and a point in B. This method avoids explicit loops and uses numpy's efficient broadcasting.