Distances from every query to every item
Medium14 min10 hidden tests
A photo app stores one embedding per photo: a vector of numbers that describes the image. When a user searches with a few query images, the app needs the Euclidean distance from every query to every stored photo. Write pairwise_distances(queries, items) that returns all of them at once, with array operations and no Python loops.
Drag the queries
d = √(|q|² + |x|² − 2 q · x)| item 0 | item 1 | item 2 | |
|---|---|---|---|
| query 0 | 0 | 10 | 3 |
| query 1 | 5 | 5 | 4 |
Query 1 to item 1: |q|² + |x|² − 2 q · x = 25 + 100 − 2 × 50 = 25, so the distance is √25 = 5. Expanded like this, rounding can push an exact 0 just below 0, so clip at 0 before the square root.
Query 1 to item 1: 5.
Examples
The query (3, 4) is 5 away from the origin and 5 away from (6, 8)
pairwise_distances([[0, 0], [3, 4]], [[0, 0], [6, 8], [3, 0]])
returns [[0, 10, 3], [5, 5, 4]]
A query that matches a stored item exactly is at distance 0
pairwise_distances([[1, 2, 2]], [[1, 2, 2], [0, 0, 0]])
returns [[0, 3]]
Requirements
querieshas shape (q, d) anditemshas shape (m, d): 2-D lists or arrays with the same number of columns d.- Return a float array of shape (q, m). Entry [i, j] is the Euclidean distance between query i and item j.
- Use array operations (broadcasting or a matrix product) instead of Python loops.
- Identical vectors are at distance 0.0. If you expand the square, rounding can make a squared distance slightly negative, so clip it at 0 before the square root.
The idea behind it
Many ML methods start from a table of distances between two sets of vectors. Nearest-neighbor search, k-means clustering, retrieval for RAG (retrieval-augmented generation) and RBF kernels all need it.
A double loop over queries and items works but is slow in Python. There are two vectorized ways. Broadcasting computes every difference at once with Q[:, None, :] - V[None, :, :], which builds a (q, m, d) array. That is fine for small inputs and wasteful for large ones. The expansion
|a - b|^2 = |a|^2 + |b|^2 - 2 a.b
needs only the squared norms of each row and one matrix product, so memory stays at (q, m). This is what libraries such as scikit-learn's euclidean_distances do, and matrix products run fast on CPUs and GPUs.
The price is rounding. Two equal vectors should give exactly 0, but the three terms can cancel to something like -1e-15, and the square root of a negative is NaN. Clipping at 0 fixes it.