Skip to content
mlmentorship

Implement KNN efficiently

The naive solution is one line. The interview is about scaling: when does naive fail, and what do you do?

Published · 7 min read ·Core ·Mixed

ML implementation · active recall

Practice before you read

40 minutes. Clarify the contract, implement a correct baseline, test edge cases, then optimize.

Back to my workbook How practice works

ML implementation · closed-book attempt

Implement KNN efficiently

Clarify the contract, implement a correct baseline, test edge cases, then optimize.

40:00recommended time

Closing or reloading clears the scratchpad. Only score, weak rubric dimensions, attempt count, and retry date can be stored locally.

30-second answer map

Visual first · depth when needed

Choose an index by separating three decisions: whether approximation is allowed, how to avoid visiting every vector, and whether stored vectors must be compressed.

Preparing the visual…

Asked in: coding round at recsys, search, and embedding-team interviews.

The L4 candidate writes the brute-force version. The L6 candidate names the regimes (small N, medium N, large N) and the right algorithm for each.

The naive solution

import numpy as np

def knn_brute(X, query, k):
    """X: (N, d) reference points. query: (d,). Returns top-k indices by L2 distance."""
    dists = np.linalg.norm(X - query, axis=1)
    return np.argpartition(dists, k)[:k]

Cost: O(N * d) per query, O(N) memory. Fine for N up to ~10K.

For batched queries, vectorize:

def knn_brute_batch(X, queries, k):
    """X: (N, d), queries: (Q, d). Returns (Q, k) of indices."""
    dists = np.linalg.norm(X[None, :, :] - queries[:, None, :], axis=2)  # (Q, N)
    return np.argpartition(dists, k, axis=1)[:, :k]

Better: use the squared-distance trick to avoid the sqrt:

# ||x - q||^2 = ||x||^2 + ||q||^2 - 2 * x.q
def knn_squared(X, queries, k):
    X_norm = (X**2).sum(axis=1)        # (N,)
    Q_norm = (queries**2).sum(axis=1)  # (Q,)
    dot = queries @ X.T                 # (Q, N)
    sq_dists = X_norm[None, :] + Q_norm[:, None] - 2 * dot
    return np.argpartition(sq_dists, k, axis=1)[:, :k]

This is the form used by FAISS’s flat index. Order-of-magnitude faster than the naive form for batched queries because of BLAS-accelerated matmul.

Learning objective: choose an index by separating three decisions: whether approximation is allowed, how to avoid visiting every vector, and whether stored vectors must be compressed.

Learning objective

Which constraint should change the KNN index?

A constraint-driven decision funnel for exact and approximate nearest-neighbor indexes Start with a flat exact scan over all N vectors and batch it on the available hardware. Only if measured latency misses its target and approximation is allowed, choose a way to visit fewer vectors: KD or ball trees for low-dimensional data where partitions prune well, HNSW graph navigation for high-recall low-latency search when graph memory is acceptable, or IVF cluster probing when scan budget is controlled by nprobe. Product quantization is a separate optional lossy compression layer when vector memory or read bandwidth is the bottleneck. Every approximate path ends by measuring recall at k against flat ground truth together with latency, memory, build time, and update cost. 1 · ESTABLISH THE EXACT BASELINE FLAT · score all N vectors exact · O(Nd) · batch / GEMM / GPU first misses latency SLO? Is approximation allowed? if no, scale or shard the flat scan 2 · IF YES, VISIT FEWER VECTORS KD / ball tree low dimension; partitions still prune effectively HNSW navigate a graph; fast, high recall; graph costs RAM IVF probe selected clusters; nprobe sets scan budget 3 · COMPRESS ONLY IF MEMORY OR READS DOMINATE OPTIONAL PQ · store short lossy codes fewer bytes · approximate distances VERIFY AGAINST FLAT GROUND TRUTH recall@k · latency · RAM · build · updates
Read it this way: start with the exact flat scan and optimize that baseline before buying index complexity. If it misses the measured latency target and approximation is acceptable, choose one candidate-pruning path for the data regime. Treat PQ as a separate compression decision; it can be combined with IVF or a graph index, but it introduces another source of distance error. Finally, compare recall@k with flat ground truth while measuring the operational costs. Original synthesis checked against the scikit-learn neighbors guide, the HNSW paper, and Faiss's research foundations.

What an L5 answer should add

“Brute-force scales as O(N * d) per query. For N up to ~100K and d ~ 100, brute-force on GPU with batched queries is fast enough.

Past that, use approximate nearest neighbor (ANN):

  • Tree methods (KD-tree, ball tree): work well for low d (up to ~20). Useless for high d (curse of dimensionality).
  • HNSW: hierarchical navigable small world graphs. State-of-the-art for moderate-d (~100-1000). Sublinear query time, high recall. Used in Qdrant, Weaviate, pgvector with HNSW.
  • IVF (inverted file): cluster the index, search only nearest clusters. Lower recall but very memory efficient. Used in FAISS.
  • PQ (product quantization): compress vectors to bytes, do approximate distance computation. Trades recall for memory dramatically.
  • HNSW + PQ: best of both worlds for billion-scale.

Trade-off knobs: recall vs query latency vs memory vs build time. Tune to your use case.”

What an L6 answer adds

“…practical things:

Distance metric matters for index choice. Cosine similarity, dot product, and L2 are different but related. For normalized vectors, cosine and L2 give the same ranking. For unnormalized, they differ. Some indexes (like inner-product HNSW) require careful normalization to be correct.

GPU vs CPU. For very high-throughput query workloads, GPU brute-force (FAISS-GPU, ScaNN) is competitive with CPU ANN up to surprisingly large N (millions). For low-throughput or memory-constrained, CPU ANN wins.

Index updates are expensive. Most ANN indexes are built once; adding/removing items is slow or unsupported. For dynamic catalogs, either rebuild periodically or use indexes with explicit incremental support (newer Qdrant, Vespa).

Recall is task-dependent. 95% recall@10 is fine for recommendation; 99% is needed for some search applications. Tune the index parameters (HNSW M / efConstruction / efSearch, IVF nprobe) to your recall floor.

Don’t roll your own. FAISS, Qdrant, ScaNN, and Annoy cover the design space. Use them. The implementation details (cache-aware data layouts, SIMD-friendly distance computation) are decade-tuned.”

Tells that get you a strong-hire vote

  • You name the squared-distance + matmul trick for batched queries.
  • You distinguish brute-force regimes from ANN regimes by N.
  • You name HNSW, IVF, PQ as the standard ANN families.
  • You mention recall vs latency vs memory trade-offs.
  • You don’t suggest implementing ANN from scratch.

Tells that get you down-leveled

  • Loops over examples in Python (use vectorized ops).
  • Computing sqrt when squared distance gives the same ranking.
  • “Use a KD-tree” for high-dimensional data.
  • No knowledge of HNSW.
  • Confusing exact KNN with ANN.

Common follow-up

“What if you have a billion vectors?”

The L6 answer:

“Single-machine memory is your floor. A billion 768-dim float32 vectors is 3 TB; doesn’t fit on one machine. Options: (1) Quantize aggressively (INT8 or PQ) to fit in memory. (2) Distribute across machines, query in parallel, merge top-K. (3) Use a hierarchical index where the top-level shards by cluster and only relevant shards are queried (IVF-style). Vespa, ScaNN, and large-scale FAISS deployments handle this. Don’t build it from scratch; the engineering is more involved than the algorithm.”

Related: Two-tower vs cross-encoder: when to use which?, Designing a RAG system that actually works, System design case study: personalized search ranking.

Found an error or stale claim? Report a correction or send a private note.