Pairwise Euclidean Distances

Compute the full matrix of Euclidean distances between two sets of points. Given X of shape (n, d) and Y of shape (m, d), produce the (n, m) matrix D where D[i, j] is the Euclidean distance between X[i] and Y[j]. This is the core of k-nearest-neighbors, k-means assignment, and RBF kernels, and interviewers use it to check whether you can vectorize a nested loop away.

Definition

The naive definition is a double loop over pairs:

D[i, j] = sqrt( sum_k (X[i, k] - Y[j, k])^2 )

The trick that removes both loops is the squared-norm expansion:

||a - b||^2 = ||a||^2 - 2 (a . b) + ||b||^2

Applied to all pairs at once, the cross term a . b becomes the single matrix product X @ Y.T, and the two norm terms broadcast across rows and columns.

Task

Complete pairwise_euclidean(X, Y) so it returns the (n, m) distance matrix, fully vectorized with NO Python loops. Use the expansion above: compute the per-row squared norms of X and Y, combine them with -2 * X @ Y.T, clip any tiny negative values (from floating-point error) up to 0, then take the square root. Do not call sklearn, scipy, or any built-in pairwise-distance helper.

Example

X = np.array([[0.0, 0.0],
              [1.0, 1.0]])
Y = np.array([[0.0, 0.0],
              [0.0, 3.0],
              [4.0, 0.0]])
pairwise_euclidean(X, Y)
# -> array([[0.        , 3.        , 4.        ],
#           [1.41421356, 2.        , 3.16227766]])

Row 0 measures the origin against each Y point; row 1 measures (1, 1) against the same three points, so D[1, 0] = sqrt(2) ≈ 1.4142.

Constraints

• 1 <= n, m <= 10^4, 1 <= d <= 10^3; use vectorized NumPy, no Python loop over pairs.
• Values fit in float64. The output is symmetric only when X and Y are the same array.
• Clip small negative squared distances to 0 before sqrt to avoid nan.

Approach

Expand ||a - b||^2 = ||a||^2 - 2 a.b + ||b||^2 and apply it to all pairs at once. Compute the per-row squared norms of X (a length-n column) and Y (a length-m row), then form the (n, m) squared-distance matrix as X_sq[:, None] - 2 * X @ Y.T + Y_sq[None, :], letting broadcasting fill the grid. Clip the result at 0 to kill any tiny negatives from floating-point cancellation, then take the square root.

Solution

import numpy as np

def pairwise_euclidean(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    Y = np.asarray(Y, dtype=np.float64)

    X_sq = np.sum(X ** 2, axis=1)        # ||X[i]||^2, shape (n,)
    Y_sq = np.sum(Y ** 2, axis=1)        # ||Y[j]||^2, shape (m,)
    cross = X @ Y.T                       # X[i] . Y[j], shape (n, m)

    # ||X[i]||^2 - 2 X[i].Y[j] + ||Y[j]||^2, via broadcasting
    sq_dist = X_sq[:, None] - 2.0 * cross + Y_sq[None, :]

    # floating-point error can push exact-zero distances slightly negative
    sq_dist = np.maximum(sq_dist, 0.0)
    return np.sqrt(sq_dist)

Walkthrough

On the example with X = [[0,0], [1,1]] and Y = [[0,0], [0,3], [4,0]]:

• X_sq = [0, 2], Y_sq = [0, 9, 16].
• cross = X @ Y.T = [[0, 0, 0], [0, 3, 4]] — for instance X[1].Y[2] = 1*4 + 1*0 = 4.
• Broadcasting X_sq[:, None] - 2*cross + Y_sq[None, :]:
  • Row 0: 0 - 2*[0,0,0] + [0,9,16] = [0, 9, 16].
  • Row 1: 2 - 2*[0,3,4] + [0,9,16] = [2, 5, 10].
• np.maximum(..., 0) leaves these unchanged, and sqrt gives [[0, 3, 4], [1.4142, 2.2360..., 3.1622...]], matching the expected output.

Complexity & notes

• Time O(nmd), dominated by the X @ Y.T matmul, which BLAS runs in optimized C. Space O(n*m) for the distance matrix plus O(n*m) for cross.
• This is much faster than the O(nmd) triple-nested Python loop and typically faster than the memory-heavy broadcasting form X[:, None, :] - Y[None, :, :], which materializes an (n, m, d) array.
• The clip before sqrt is essential: when a point in X equals a point in Y, the squared distance is mathematically 0 but can come out as a small negative (e.g. -1e-15), and np.sqrt of a negative is nan.
• Casting to float64 guards against integer inputs (where X ** 2 could overflow or truncate) and against float32 inputs where the cancellation error is larger.
• To recover distances within a single set, call pairwise_euclidean(X, X); the diagonal will be exactly 0 thanks to the clip.