Nearest Image Embeddings
MediumArrays · Apple
An image search service stores a gallery of image embeddings, each an integer vector, and receives a list of queries, also integer vectors. All vectors have the same dimension.
The distance between two vectors is their squared Euclidean distance: the sum of (a[d] - b[d])^2 over every dimension d. (No square root is needed, so everything stays an integer.)
For each query, return the indices of the k gallery vectors closest to it, nearest first. If two gallery vectors are equally close, the smaller index comes first. If the gallery has fewer than k vectors, return all of them in that order.
Return one list per query, in the same order as queries.
Example 1
- Input:
- gallery = [[0, 0], [3, 4], [1, 1], [-2, 0]], queries = [[2, 2], [-1, 0]], k = 2
- Output:
- [[2, 1], [0, 3]]
- Explanation
- For query
[2, 2]the squared distances are 8, 5, 2 and 20, so the two nearest are index 2 then index 1. For query[-1, 0]indices 0 and 3 are both at distance 1, so the tie keeps index 0 first.
Example 2
- Input:
- gallery = [[5], [1], [3]], queries = [[2]], k = 3
- Output:
- [[1, 2, 0]]
- Explanation
- The distances from 2 are 9, 1 and 1. Indices 1 and 2 tie and are listed by index, then index 0.
Example 3
- Input:
- gallery = [[1, 2, 3]], queries = [[0, 0, 0]], k = 4
- Output:
- [[0]]
- Explanation
- The gallery has only one vector, so even though
kis 4 the result holds just index 0.
Constraints
1 <= gallery.length <= 20000 <= queries.length <= 2001 <= dimension <= 64-1000 <= vector values <= 10001 <= k <= 50