Your question is Efficient k-NN for Embeddings. Start with the requirements on the right.
Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.
Arrowstreet Capital's research systems compare high-dimensional embeddings repeatedly. Implement an exact k-nearest neighbors search that avoids sorting every distance for each query by indexing the embeddings in a vantage-point tree.
Write knn_search(embeddings, query, k). embeddings is a non-empty list of equal-length numeric vectors, and query is a vector with the same dimension. Use squared Euclidean distance. Return the indices of the k closest embeddings, ordered by increasing distance. Break equal-distance ties by increasing original index.
The function may construct the VP-tree during the call. The returned result must be exact, not approximate. A VP-tree recursively selects a pivot, partitions points by their distance from that pivot, and uses triangle-inequality bounds to prune subtrees during search. Maintain only the best k candidates while traversing.
Example 1
Input: embeddings = [[0, 0], [2, 0], [0, 3], [-1, 0], [5, 5]], query = [0, 0], k = 3
Output: [0, 3, 1]
Distances are 0, 1, 4, 9, and 50, respectively.
Example 2
Input: embeddings = [[1, 1], [1, 1], [2, 2]], query = [1, 1], k = 2
Output: [0, 1]
The first two vectors are tied at distance zero, so their original indices determine the order.
1 <= len(embeddings) <= 100,0002 <= len(embeddings[0]) <= 5121 <= k <= len(embeddings)def knn_search(embeddings, query, k):