Vector databases retrieve embeddings that point in a direction similar to a query. Cosine similarity compares direction while normalizing away magnitude.
Design a CosineVectorSearch class:
CosineVectorSearch() creates a stateless exact-search helper.int[] topK(int[] query, int[][] vectors, int k) returns the indices of the k most similar stored vectors.
For vectors a and b:
Rank by cosine similarity descending. When two scores are equal, place the smaller vector index first.
Example 1:
Input:
Output:
Explanation: Index 0 has cosine 1. Index 2 has cosine about 0.707, ahead of the orthogonal and opposite vectors.
Example 2:
Input:
Output:
Explanation: Indices 0 and 1 point in the same direction as the query and tie, so index order decides between them. Indices 2 and 3 also tie, making 2 the third result.
Constraints
1 <= query.length <= 5001 <= vectors.length <= 10^4vectors[i].length == query.length1 <= k <= vectors.length- Coordinates are integers from
-10^4 to 10^4. - Neither the query nor any stored vector is the zero vector.
- At most
100 calls are made to topK.