Top-k by cosine similarity

Every vector database answers one question: which stored vectors point the same way as this one? Build the brute-force version that every index is measured against.

Write top_k(query, docs, k):

• query is a list of floats.
• docs is a list of vectors, each the same length as query.
• Return the indices of the k docs with the highest cosine similarity to query, best first.
• Break ties by the lower index.
• A zero vector has similarity 0 with everything (don’t divide by zero).
• If k is larger than len(docs), return every index.

Cosine similarity is the dot product divided by the product of the lengths:

cos(a, b) = (a · b) / (|a| × |b|)

No numpy. Plain Python only.

>>> top_k([1, 0], [[0, 1], [1, 1], [2, 0]], 2)
[2, 1]

Approach

Score every document, then sort indices by score (descending) with the index as the tie-breaker.

import math

def top_k(query, docs, k):
    def norm(v):
        return math.sqrt(sum(x * x for x in v))

    qn = norm(query)
    scores = []
    for d in docs:
        dn = norm(d)
        dot = sum(a * b for a, b in zip(query, d))
        scores.append(0.0 if qn == 0 or dn == 0 else dot / (qn * dn))
    return sorted(range(len(docs)), key=lambda i: (-scores[i], i))[:k]

Complexity

O(n·d) to score n documents of dimension d, plus O(n log n) to sort. Swap the sort for heapq.nsmallest(k, …) to get O(n log k) when k is small.

Why this matters

This is exact search. Approximate indexes (HNSW, IVF) trade a little recall for a lot of speed, and you measure that recall against exactly this function. Precompute qn once: it’s the same for every document, and at a million documents that saves a million square roots.