1007. Top-K Retrieval by Cosine Similarity

Medium · Math

You have N stored vectors and a query vector. Return the indices of the K most similar stored vectors, sorted by cosine similarity (descending).

Foundation of vector-database lookup; every RAG pipeline does this.

Input (stdin, JSON): { "query": [1, 0, 0], "vectors": [[1, 0, 0], [0, 1, 0], [0.7, 0.7, 0]], "k": 2 }

Output (stdout): JSON list of K indices, e.g. "[0, 2]".

Tie-break by lower index first.

Examples

Example 1
Input: {"query":[1,0,0],"vectors":[[1,0,0],[0,1,0],[0.7,0.7,0]],"k":2}
Output: [0, 2]
Explanation: Vector 0 is identical (sim=1), vector 2 is 45° off (sim≈0.707).

Constraints