Cosine Similarity Top-k Search in NumPy

Given a query q (length d, not all zero) and a matrix M (n × d), return the indices of the k rows with the highest cosine similarity to q, most similar first.

A row that is entirely zero has similarity 0.
Similarities are compared after rounding to 12 decimal places; equal values are ordered by row index, lower first.

Examples

Input q = [1, 0], M = [[1, 0], [0, 1], [1, 1], [-1, 0]], k = 2

Output [0, 2]

Row 0 points the same way (similarity 1), row 2 at 45 degrees.

Input q = [1, 1], M = [[2, 2], [1, 1], [0, 0], [3, 3]], k = 3

Output [0, 1, 3]

Rows 0, 1 and 3 have the same direction: a three-way tie, broken by index.

Constraints

1 ≤ k ≤ n ≤ 10⁴, 1 ≤ d ≤ 256

Hints

Hint 1

Divide M by its row norms, replacing zero norms with 1 so that zero rows stay zero.

Hint 2

One matrix-vector product gives all n similarities; divide by the query's norm.

Hint 3

np.lexsort((index, -rounded)) sorts by similarity descending, then by index.

Where candidates go wrong

Dividing by a zero norm: nan similarities that sort unpredictably.
Comparing raw floats for ties: vectors with the same direction can differ in the last bit after normalisation.
Returning the k largest dot products, which favours long vectors.

The idea

Normalise the rows once (zero rows stay zero), take one matrix-vector product with the normalised query, round to 12 places, and sort by (−similarity, index). For large n, np.argpartition finds the top k in linear time before sorting just those k.

Time O(nd + n log n) · Space O(n)

Reference solution

import numpy as np

def cosine_top_k(q, M, k):
    q = np.asarray(q, dtype=float)
    M = np.asarray(M, dtype=float)
    norms = np.linalg.norm(M, axis=1)
    safe = np.where(norms == 0, 1.0, norms)            # zero rows stay zero
    sims = (M / safe[:, None]) @ (q / np.linalg.norm(q))
    key = np.round(sims, 12)
    order = np.lexsort((np.arange(len(M)), -key))      # by similarity desc, then index
    return [int(i) for i in order[:k]]

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

The core of every vector database and every retrieval step in a RAG system. The interview checks vectorisation, the zero-vector case and a reproducible tie rule.

The theory behind it is in the lesson Cosine search, contrastive loss and negative sampling.

Keep going

Easy12 min

Top-k search by cosine similarity