IVF (Inverted File) Search

Background

IVF (Inverted File Index) clusters the database into Voronoi cells around K-means centroids. At query time, you only search the top-nprobe cells — quadratic-to-linear speedup in exchange for some recall. The simple ANN baseline used by FAISS at billion-scale.

Problem statement

Implement ivf_search(query, db, centroids, assignments, k, nprobe=2). Steps:

1. Compute the L2 distance from query to each centroid.
2. Pick the nprobe nearest centroids.
3. Restrict candidates to db rows whose assignments lie in those cells.
4. Within the candidates, compute exact L2 distance to the query.
5. Return the top-$k$ candidate indices.

Input

• query — 1-D np.ndarray of length $d$.
• db — 2-D np.ndarray shape $(n, d)$.
• centroids — 2-D np.ndarray shape $(C, d)$.
• assignments — 1-D np.ndarray of int length $n$, the centroid index for each db row.
• k — int, target result count.
• nprobe — int >= 1, number of cells to probe (default $2$).

Output

Returns list[int] of length $\min(k, |\text{candidates}|)$ — db indices in ascending-distance order.

Examples

Example 1 — query inside one cell

Input:  query close to centroid 0; nprobe=1
Output: indices of the rows assigned to centroid 0, ordered by distance

Example 2 — nprobe spans multiple cells

nprobe=2 returns the union of candidates from the two nearest cells.

Constraints

• If no db row falls in any probed cell, return [].
• Returns Python ints, not numpy ints.
• k may exceed candidate count; in that case return all candidates.

Notes

• Recall vs latency. nprobe = 1 is fastest but lossy. Doubling nprobe roughly halves the recall miss rate.
• Train the centroids on a subsample. Production IVF clusters on 1M+ vectors with K-means; the centroids here are an input, not learned.