Problem 109 · Retrieval

Cosine Top-k with a Bounded Heap

MediumAI engineering

01Problem

Write `top_k_cosine(query, corpus, k)` where `corpus` maps doc id to an embedding. Return the `k` most similar `(id, score)` pairs by cosine similarity, highest first, ties broken by smaller id. Skip zero vectors (cosine is undefined), raise on a zero query or a dimension mismatch, and use a size-k min-heap so memory stays O(k) for any corpus size.

02Key insight

Normalise once and cosine becomes a dot product; in production you store unit vectors so the index can use inner product. A size-k min-heap whose root is the current k-th best lets you reject most candidates with one comparison, giving O(n log k) instead of sorting all n. Zero vectors are a real data-quality signal (empty chunks, failed embedding calls) and must not silently score 0 and crowd the results. This brute-force scan is also the ground truth you use to measure an ANN index's recall.

03Solution

solution.pypython43 lines

import heapq
import math


def _unit(v: list[float]) -> list[float] | None:
    n = math.sqrt(sum(x * x for x in v))
    return None if n == 0 else [x / n for x in v]


def top_k_cosine(query: list[float], corpus: dict[str, list[float]], k: int) -> list[tuple[str, float]]:
    q = _unit(query)
    if q is None:
        raise ValueError("zero query vector")
    if k <= 0:
        return []
    heap: list[tuple[float, str]] = []  # min-heap keyed by (score, inverted id order)
    for doc_id, vec in corpus.items():
        if len(vec) != len(q):
            raise ValueError(f"dimension mismatch for {doc_id}")
        u = _unit(vec)
        if u is None:
            continue
        score = sum(a * b for a, b in zip(q, u))
        item = (score, _Desc(doc_id))
        if len(heap) < k:
            heapq.heappush(heap, item)
        elif item > heap[0]:
            heapq.heapreplace(heap, item)
    return [(d.value, s) for s, d in sorted(heap, reverse=True)]


class _Desc:
    """Wraps an id so that a smaller id compares as larger (wins ties in a max ordering)."""
    __slots__ = ("value",)

    def __init__(self, value: str) -> None:
        self.value = value

    def __lt__(self, other: "_Desc") -> bool:
        return self.value > other.value

    def __eq__(self, other: object) -> bool:
        return isinstance(other, _Desc) and self.value == other.value

Try it first.

TimeO(n * d + n log k)

SpaceO(k)

tests.py · 7 passing

tests.pypython20 lines

import random
corpus = {"a": [1.0, 0.0], "b": [0.0, 1.0], "c": [1.0, 1.0], "z": [0.0, 0.0], "d": [2.0, 0.0]}
res = top_k_cosine([1.0, 0.1], corpus, 2)
assert [d for d, _ in res] == ["a", "d"]
assert abs(res[0][1] - res[1][1]) < 1e-12
assert top_k_cosine([1.0, 0.0], corpus, 10)[-1][0] == "b"
assert len(top_k_cosine([1.0, 0.0], corpus, 10)) == 4
assert top_k_cosine([1.0, 0.0], corpus, 0) == []
for bad in ([0.0, 0.0], [1.0, 0.0, 0.0]):
    try:
        top_k_cosine(bad, corpus, 1)
        assert False
    except ValueError:
        pass
rng = random.Random(1)
big = {f"d{i:03d}": [rng.gauss(0, 1) for _ in range(8)] for i in range(300)}
q = [rng.gauss(0, 1) for _ in range(8)]
cos = lambda u, v: sum(a * b for a, b in zip(u, v)) / math.sqrt(sum(a * a for a in u) * sum(b * b for b in v))
brute = sorted(big, key=lambda d: (-cos(q, big[d]), d))[:10]
assert [d for d, _ in top_k_cosine(q, big, 10)] == brute

04Staff-level follow-up

At 100M vectors brute force is too slow, so you move to HNSW or IVF-PQ and measure recall@k against exactly this function on a held-out query set. Discuss filtered search (tenant or ACL filters) and why post-filtering an ANN top-k can return fewer than k results; pre-filtered or partitioned indexes fix it.