Given a query vector and N stored vectors, return the top-k most similar by cosine similarity, efficiently.

TL;DR: Cosine similarity is the dot product of L2-normalized vectors. Normalize once, compute all similarities as a single matrix-vector product (vectorized, not a Python loop), then take the top-k with a partial selection (np.argpartition, O(N)) rather than a full O(N log N) sort. This is brute-force exact search, O(N·d); at large N you switch to an ANN index.

How to approach it. State that cosine similarity = normalized dot product, so pre-normalizing turns the whole thing into one matrix multiply. Then note the top-k should use partial selection, not a full sort. Write it vectorized, and raise the scaling follow-up (ANN) since brute force is O(N) per query.

A strong answer.

import numpy as np

def top_k_cosine(query, vectors, k):
    # query: (d,), vectors: (N, d)
    q = query / (np.linalg.norm(query) + 1e-12)
    V = vectors / (np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-12)
    sims = V @ q                                   # (N,) cosine sims, one matmul
    # partial selection: top-k indices in O(N), not a full sort
    idx = np.argpartition(-sims, kth=k-1)[:k]
    idx = idx[np.argsort(-sims[idx])]              # sort just the k winners
    return list(zip(idx.tolist(), sims[idx].tolist()))

What a strong candidate explains:

• Cosine = normalized dot product. cos(a,b) = (a·b)/(‖a‖‖b‖). Normalize each vector to unit length once, and similarity is just the dot product, so all N similarities are a single matrix-vector product V @ q, fully vectorized (no Python loop over rows). The 1e-12 guards against a zero-norm vector.
• Partial selection beats full sort. You want the top k, not a fully sorted list, so np.argpartition finds the k largest in O(N) and you sort only those k (O(k log k)), instead of O(N log N) to sort everything. On large N this matters.
• Pre-normalize the corpus once. If you query repeatedly, normalize vectors at index-build time so each query is just V @ q; do not renormalize every call.
• Numerical note. Normalizing makes cosine equivalent to ranking by dot product, and (for unit vectors) by Euclidean distance, useful because some indexes use inner-product or L2 metrics.

Key takeaways

• Cosine on L2-normalized vectors is just a dot product, so all N similarities collapse to one V @ q matmul.
• np.argpartition selects the top-k in O(N); sort only those k. A full argsort is wasteful O(N log N).
• Pre-normalize the corpus at index-build time so each query skips renormalization.
• Brute force is O(N·d) per query; past millions of vectors switch to an ANN index (HNSW/IVF) and trade a little recall for speed.

What interviewers probe next.

• "Complexity, and what about millions of vectors?" This is O(N·d) per query (brute-force exact); it does not scale to millions at high QPS, so use an ANN index (HNSW/IVF), trading a little recall for speed.
• "Cosine vs dot product vs Euclidean?" On normalized vectors they rank consistently; cosine ignores magnitude (good for embeddings where direction is meaning), raw dot product does not.
• "Why argpartition over argsort?" O(N) partial selection vs O(N log N) full sort; you only need the k best, unordered, then sort those k.
• "Batch many queries?" Stack queries into a matrix Q (m, d) and compute Q @ V.T for all pairwise sims in one matmul, then top-k per row.

Common mistakes.

• Looping over vectors in Python instead of one vectorized matmul.
• Forgetting to normalize (then you compute dot product, not cosine) or dividing by a zero norm.
• Doing a full sort when a partial top-k selection is O(N).
• Claiming brute force scales; at large N you need an approximate index.