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 · 8 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.

Problem

Context. Reference vectors X have shape (N, d); a query has shape (d,), or batched queries have shape (Q, d). Distance is Euclidean distance over matching numeric features. Task. Return k original reference-row indices per query with the smallest distances, then compare exact and approximate scaling options. For reference points [[0, 0], [3, 0], [8, 0]], query [2, 0], and k = 1, return index 1. Constraints. The provided partition-based sketches assume finite, representable arithmetic and 1 <= k < N. The selected indices need not be distance-sorted; ties may select any valid neighbors. Preserve inputs. Approximate search requires an explicit recall target and comparison against exact ground truth.

40:00recommended time

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

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

Find the nearest reference vectors to each query and explain how the search should scale as the collection grows.

Problem

Context. Reference vectors X have shape (N, d); a query has shape (d,), or batched queries have shape (Q, d). Distance is Euclidean distance over matching numeric features.

Task. Return k original reference-row indices per query with the smallest distances, then compare exact and approximate scaling options. For reference points [[0, 0], [3, 0], [8, 0]], query [2, 0], and k = 1, return index 1.

Constraints. The provided partition-based sketches assume finite, representable arithmetic and 1 <= k < N. The selected indices need not be distance-sorted; ties may select any valid neighbors. Preserve inputs. Approximate search requires an explicit recall target and comparison against exact ground truth.

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

IVF scans fewer vectors; PQ stores fewer bytes per vector

Scanning fewer vectors and storing fewer bytes are different savingsIllustrative vector blocks distinguish read count from representation size. IVF scans fewer vectors; PQ stores fewer bytes per vector. Neither recovers a candidate excluded by the scan decision.Read count and bytes per vector aredifferent dimensionsBaselineIVF: fewer read vectorsPQ: smaller vector representationBoth dimensions reduced
Derivation
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
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.
Check: Which choice reduces visited vectors, which reduces bytes per vector? What exact baseline measures the lost recall?
Read it this way: Illustrative vector blocks distinguish read count from representation size. IVF scans fewer vectors; PQ stores fewer bytes per vector. Neither recovers a candidate excluded by the scan decision.
Which change saves representation bytes without reducing scanned-vector count?

What an L5 answer should add

“Brute-force distance evaluation costs O(N * d) per query. Whether it meets the budget depends on hardware, batching, distance, precision, and data movement; a row-count threshold is not a latency guarantee.

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.
  • Graph search plus quantization: another candidate with build, recall, memory, and update trade-offs; no combination is universally best at a catalog size.

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

What an L6 answer adds

“…practical things:

Preserve the intended score. Nonzero unit-normalized vectors have equivalent cosine, dot-product, and squared-L2 rankings. Normalizing arbitrary item vectors can change an inner-product task; inner-product search does not universally require it. Match query/index transforms and test against an exact top-k oracle with a declared tie policy.

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.

Updates are an index contract. Insertion, deletion, compaction, and concurrent-query behavior vary by implementation. For a changing catalog, evaluate freshness and update cost alongside query recall instead of assuming every index is static.

Recall requirements are task-dependent. Set an acceptable miss cost and recall floor from the application; neither a recommendation nor search label determines a universal percentage. Tune index parameters against that floor and measured serving cost.

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.