Top-K Retrieval

Start from brute force

The obvious approach, and the one the reference takes, scores all n documents, sorts all n by similarity, and keeps the first k. It is correct: a sort key of (-similarity, index) gives the required order, ties included, and slicing past the end is safe. It costs O(n·d + n log n) time for d-dimensional vectors and O(n) space for the scored list. The sort orders n − k documents that are immediately thrown away, which matters when k is 5 and n is in the millions. Even then, the larger cost is the O(n·d) scoring pass, which only an index can avoid.

The approach

Score once, then select without a full sort. Keep a min-heap of size k ordered worst-first: with Python’s heapq, push (similarity, -index), so the root has the lowest similarity and, among equals, the larger index; whenever the heap exceeds k entries, pop the root. After the scan the heap holds exactly the top k, and sorting those k gives the output order. heapq.nsmallest(k, ...) on (-similarity, index) pairs does the same and is documented as equivalent to sorted(...)[:k].

The reference keeps the full sort for clarity. It builds scored as (cosine_similarity(query, d), i) tuples, sorts with key=lambda t: (-t[0], t[1]), which ranks by similarity descending and then index ascending, and returns the indices from scored[:k]. Cosine ignores vector length, so [2, 0] and [1, 0] score identically against [1, 0].

Walk through an example

On the first example, [1, 1] is 45 degrees from the query and scores √2/2, enough to beat the orthogonal [0, 1] for the second slot.

Complexity

Scoring is O(n·d) for n documents of dimension d. The reference’s full sort adds O(n log n) time and keeps O(n) space for scored; a size-k heap makes selection O(n log k) with O(k) extra space, and quickselect is O(n) on average. At millions of vectors the O(n·d) scan dominates, which an approximate index such as HNSW avoids.

Edge cases to test

• Ties break by index, and cosine ignores length: with q = [1, 0], the docs [2, 0] at index 1 and [1, 0] at index 4 both score exactly 1.0, so index 1 comes first.
• k beyond the list: top_k([1, 0], [[0, 1], [1, 0]], 5) returns [1, 0], because slicing past the end is safe, and k = 0 returns [].
• A negative k is not rejected: scored[:-1] and slice(0, -1) both mean “all but the last”, so k = -1 on the first example’s three docs silently returns [0, 2]. Validate k ≥ 0.
• Opposite vectors score −1 and still rank, last: top_k([1, 0], [[-1, 0], [0, 1]], 2) returns [1, 0]. Add a similarity threshold if weak or negative matches should not be returned at all.
• A zero vector has no direction. Python’s cosine_similarity raises ZeroDivisionError. TypeScript gets NaN, and since NaN || (a.i - b.i) falls through to the index comparison, a zero vector at index 0 sorts ahead of a perfect match: topK([1,0], [[0,0],[1,0]], 1) returns [0]. Reject zero vectors at ingestion.

Common mistakes

• Sorting all n scores at scale when only k are needed: fine for a few hundred documents, wasteful for millions.
• Using a max-heap of all n elements instead of a bounded min-heap of k.
• Mishandling ties or fewer than k documents.
• Sorting the (similarity, index) tuples with reverse=True. That reverses the tie-break too, so among equal scores the larger index comes first.

Follow-ups an interviewer asks next

The corpus holds 10 million vectors. How do you serve queries fast?

Use an approximate nearest-neighbor index, such as HNSW or IVF with product quantization, which examines a small fraction of the vectors per query and still returns most of the true top k. Measure recall@k against this exact brute force on sample queries. Normalize vectors at ingestion so cosine becomes a dot product, which vector indexes compute natively.

How do you combine this with keyword search?

Run BM25 alongside it and merge the two ranked lists, for example with Reciprocal Rank Fusion, which gives each document the sum of 1 / (60 + rank) across the lists. Keyword search catches exact identifiers and rare names that embeddings blur, and rank fusion needs no calibration between the two scoring scales.

The top k are relevant but redundant. What next?

Rerank the candidates. A cross-encoder rescoring the top 50 to 100 improves precision, and Maximal Marginal Relevance trades a little relevance for diversity so near-duplicate chunks do not fill the prompt. Both run on a short list, so their extra cost per query stays small.

Can you avoid recomputing norms on every query?

Yes. Store each document vector already divided by its norm. Cosine with a normalized query is then a plain dot product, and scoring every document is one matrix-vector multiply. It also takes the zero-vector division out of the query path, since zero vectors are rejected once, at ingestion.

Where this shows up

This is the “R” in RAG: rank documents by embedding similarity and take the top k to put in the prompt. Retrieval quality caps answer quality, and at scale this exact top-k step is where an approximate-nearest-neighbor index (HNSW) replaces the brute-force scan to keep it fast.

The brute-force version stays useful even after an index arrives. It is the ground truth for measuring an approximate index: run both on a sample of real queries and report recall@k, the fraction of the exact top k that the index also returned. It is also the right production choice for small corpora, such as one user’s notes or one tenant’s documents, where scanning a few thousand normalized vectors is one matrix-vector multiply and an index would only add build time, memory and a recall loss. The same select-k pattern appears in recommendation candidate generation, in deduplication (finding near-identical chunks before indexing them) and in classic k-nearest-neighbor classification, which differ mainly in what the vectors encode and what happens to the k results.

Check yourself

Reference solution

import math

def cosine_similarity(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)


def top_k(query, docs, k):
    # score every doc, then rank by similarity (desc), tie-break by index (asc)
    scored = [(cosine_similarity(query, d), i) for i, d in enumerate(docs)]
    scored.sort(key=lambda t: (-t[0], t[1]))
    return [i for _, i in scored[:k]]


print(top_k([1, 0], [[1, 0], [0, 1], [1, 1]], 2))   # -> [0, 2]

function cosineSimilarity(a: number[], b: number[]): number {
  let dot = 0, na = 0, nb = 0;
  for (let i = 0; i < a.length; i++) { dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; }
  return dot / (Math.sqrt(na) * Math.sqrt(nb));
}

function topK(query: number[], docs: number[][], k: number): number[] {
  // score every doc, then rank by similarity (desc), tie-break by index (asc)
  const scored = docs.map((d, i) => ({ s: cosineSimilarity(query, d), i }));
  scored.sort((a, b) => (b.s - a.s) || (a.i - b.i));
  return scored.slice(0, k).map((x) => x.i);
}

console.log(topK([1, 0], [[1, 0], [0, 1], [1, 1]], 2));   // -> [0, 2]