Pairwise Squared Distances

Compute the squared distance from every point to every center without Python loops.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Pairwise Squared Distances: Broadcast every point-center pair across a feature axis, square coordinate differences, then reduce that axis.

Input and goalCompute the squared distance from every point to every center without Python loops.

Lay out points and centersUse points P0=(0,0), P1=(1,2) and centers C0=(1,0), C1=(2,2). The output must contain all four P-C pairs.

P0 (0,0)P0 (0,0)P1 (1,2)P1 (1,2)

columnsC0 (1,0), C1 (2,2)shapepair grid [n=2,k=2]

Recognize it

Use it when every item in one batch must pair with every item in another batch and the per-pair calculation is identical across a shared feature dimension.

Keep true

Cell [i,j,:] always belongs to point i and center j; broadcasting changes alignment, not values, and reducing only the final feature axis preserves the pair grid.

Reuse it

For all-pairs tensor operations, insert singleton axes where each operand should repeat, verify the broadcasted axis meaning, and reduce only the intended feature axes.

Pattern: Broadcasting.

Simple idea: Change points from shape (n, d) to (n, 1, d) and centers from (k, d) to (1, k, d). Their difference has shape (n, k, d). Sum over the last axis.

import numpy as np

def pairwise_squared_distances(
   points: np.ndarray, centers: np.ndarray
) -> np.ndarray:
   differences = points[:, None, :] - centers[None, :, :]
   return np.sum(differences * differences, axis=-1)

Cost: $O(nkd)$ time and $O(nkd)$ temporary space.