AlgoMaster Logo
AlgoMasterFind Nearest Vectors by Cosine Similaritymedium

Find Nearest Vectors by Cosine Similarity

medium

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 <= 500
  • 1 <= vectors.length <= 10^4
  • vectors[i].length == query.length
  • 1 <= 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.
Hints

Loading...
CallReturns
new CosineVectorSearch()null
topK([1,0], [[1,0],[0,1],[1,1],[-1,0]], 2)[0,2]

The scores are 1, 0, about 0.707, and -1, so indices 0 and 2 rank first.

Run checks these cases. Submit also runs a larger hidden set.