Givzey's embedding retrieval pipeline stores candidate embeddings sorted by their first coordinate. Given a sorted query embedding array queries and a sorted candidate embedding array candidates, return the indices of the k nearest candidates for every query using squared Euclidean distance.
For each query, use the ordering to avoid scanning candidates that cannot improve the current top-k results. Return one list per query, with candidate indices ordered by increasing distance. If distances tie, any ordering among tied candidates is valid.
Implement k_nearest_neighbors(queries, candidates, k).
queries is a list of n embeddings, each represented by a list of d numbers.candidates is a list of m embeddings, each represented by a list of d numbers.candidates is sorted in nondecreasing order by embedding[0].queries is also sorted by embedding[0], although each query must still be processed independently.n lists. Each inner list contains the indices of the k nearest candidates.def k_nearest_neighbors(queries, candidates, k):