Cosine Similarity Matrix

Compute the full matrix of cosine similarities between the rows of a data matrix. Given X of shape (n, d), produce the (n, n) matrix S where S[i, j] is the cosine similarity between rows X[i] and X[j]. Cosine similarity is the workhorse of text and embedding retrieval, and interviewers use it to check whether you can turn a per-pair formula into one clean normalize-then-multiply.

Definition

Cosine similarity measures the angle between two vectors, ignoring their magnitudes:

S[i, j] = (X[i] . X[j]) / (||X[i]|| * ||X[j]||)

If you first scale each row to unit length, u = x / ||x||, then the formula collapses to a plain dot product:

S[i, j] = u_i . u_j

So normalizing every row once and taking X_norm @ X_norm.T produces the whole matrix in a single matrix product.

Task

Complete cosine_similarity_matrix(X) so it returns the (n, n) similarity matrix, fully vectorized with NO Python loops. Compute each row’s L2 norm, divide each row by its norm to get unit vectors, then multiply the unit matrix by its transpose. A zero-norm row (all zeros) has no direction: treat its unit vector as all zeros so its similarity with everything, including itself, is 0 rather than nan. Do not call sklearn, scipy, or any built-in cosine helper.

Example

X = np.array([[1.0, 0.0],
              [0.0, 2.0],
              [1.0, 1.0],
              [0.0, 0.0]])
cosine_similarity_matrix(X)
# -> array([[1.        , 0.        , 0.70710678, 0.        ],
#           [0.        , 1.        , 0.70710678, 0.        ],
#           [0.70710678, 0.70710678, 1.        , 0.        ],
#           [0.        , 0.        , 0.        , 0.        ]])

Row 0 is orthogonal to row 1 (similarity 0) and at 45 degrees to row 2 (similarity 1/sqrt(2) ≈ 0.7071). Row 3 is the zero vector, so every entry in its row and column is 0.

Constraints

• 1 <= n <= 10^4, 1 <= d <= 10^3; use vectorized NumPy, no Python loop over pairs.
• Values fit in float64. The output is symmetric and each non-zero row has S[i, i] = 1.
• Handle zero-norm rows without dividing by zero: their similarities are 0, not nan.

Approach

Cosine similarity is the dot product of unit-normalized vectors, so normalize once and multiply. Compute each row’s L2 norm, replace any zero norm with 1 before dividing so zero rows stay all-zero instead of becoming nan, then form the unit matrix X_norm and return X_norm @ X_norm.T. Because every retained row is unit length, the matrix product delivers all pairwise similarities at once.

Solution

import numpy as np

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

    norms = np.sqrt(np.sum(X ** 2, axis=1, keepdims=True))  # ||X[i]||, shape (n, 1)

    # avoid divide-by-zero: a zero row divided by 1 stays all zeros
    safe = np.where(norms == 0.0, 1.0, norms)
    X_norm = X / safe                                       # unit rows, shape (n, d)

    return X_norm @ X_norm.T                                # S[i, j] = u_i . u_j

Walkthrough

On the example with rows [1,0], [0,2], [1,1], [0,0]:

• norms = [[1], [2], [sqrt(2)], [0]].
• safe swaps the last 0 for 1, giving [[1], [2], [sqrt(2)], [1]].
• X_norm = [[1,0], [0,1], [0.7071,0.7071], [0,0]] — each non-zero row is now unit length, and the zero row stays zero.
• X_norm @ X_norm.T: row 0 dotted with row 2 is 1*0.7071 + 0*0.7071 = 0.7071; row 0 with row 1 is 0; every dot with row 3 is 0 because it is all zeros. The diagonal is 1 for non-zero rows and 0 for the zero row, matching the expected output.

Complexity & notes

• Time O(n^2 * d), dominated by the X_norm @ X_norm.T matmul that BLAS runs in optimized C. Space O(n^2) for the output plus O(n*d) for the normalized copy.
• The zero-norm guard is the whole trick: dividing a zero row by its 0 norm gives nan, which then poisons an entire row and column of the product. Substituting 1 in the denominator leaves the zero row untouched, so its similarities come out as an honest 0.
• Normalizing once (O(n*d)) and multiplying is far cheaper than computing n^2 separate norm products; the naive per-pair form recomputes ||X[i]|| * ||X[j]|| for every pair.
• Casting to float64 guards against integer inputs where X ** 2 could overflow, and keeps the normalization accurate for float32 data.
• Floating-point rounding can push a diagonal or identical-row entry a hair above 1.0; clip to [-1, 1] if you later feed the result into arccos to recover angles.