Search an IVF index
centroids are the cluster centres, lists[c] holds the ids of the points in cluster c, and points[i] is point i (all 2-D (x, y) pairs). dist2(a, b) is given.
Write ivf_search(query, points, centroids, lists, nprobe, k) that opens the nprobe clusters whose centroids are closest to query, ranks the points in them by squared distance (ties: the point met first in the opened lists wins) and returns a pair (ids, computed): the best k ids, and the number of distances computed, which is one per centroid plus one per point in the opened lists.
What Check tests
- nprobe 1 searches only the nearest cluster
- nprobe 1 misses the true nearest point
- nprobe 2 finds it, at a higher cost
- The cost counts the centroids
- k larger than the candidates returns what exists
Stuck? One hint at a time
The solution unlocks after your first Check. Try first: that is where the learning happens.
Output
Press Run to see what your program prints.