Implementationmedium~30 min
Objective
Build two vector search indices (exact brute-force and approximate IVF with k-means clustering), then compare their speed/recall tradeoffs.
Background
Your team is building a semantic search feature for a document store with 100K+ embeddings. Brute-force cosine similarity is correct but too slow at scale. You need an approximate nearest neighbor (ANN) index that trades a small amount of recall for a large speedup. Implement both a brute-force baseline and an IVF index from scratch using only numpy, then benchmark them to understand the speed-recall tradeoff.
Requirements
- 1.Implement brute-force search using cosine similarity with proper normalization
- 2.Implement k-means clustering for the IVF training step
- 3.Implement IVF add: assign vectors to their nearest centroid bucket
- 4.Implement IVF search: probe n_probe nearest centroids, search within those buckets
- 5.Return results as (distances, ids) tuples with correct top-k selection
Evaluation (100 points)
Brute force search with cosine similarity
Uses dot product / norm for similarity and sorts to find top-k
K-means clustering for IVF training
Implements centroid initialization, assignment, and mean update loop
IVF add assigns vectors to buckets
Computes nearest centroid for each vector and stores in self.buckets
IVF search probes nearest centroids
Selects n_probe nearest centroids and searches within those buckets
Returns (distances, ids) tuples
Search methods return properly formatted (distances, ids) pairs