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):