Vector Search from Scratch | Dataford Interview Questions

PRGX Global's spend analytics workflows can represent records as numeric vectors. Implement a basic vector search function that ranks candidate vectors by cosine similarity to a query vector and returns the indices of the top k matches.

Formal Specification

Given vectors, a list of n equal-length numeric vectors, a nonzero query vector, and an integer k, return a list of the k vector indices with the highest cosine similarity. Order results by decreasing similarity. If two vectors have the same similarity, place the smaller index first.

Compute similarity as:

cosine(a, b) = dot(a, b) / (magnitude(a) * magnitude(b))

Return only indices, not similarity scores. Do not use machine learning or vector-search libraries. A heap-based solution is encouraged so that retaining the best matches does not require sorting every vector.

Constraints

1 <= len(vectors) <= 10^5
1 <= len(query) <= 100
Every vector has the same dimension as query
Every vector and the query have nonzero magnitude
1 <= k <= len(vectors)
Vector components are in [-10^4, 10^4]

Function Signature

def vector_search(vectors, query, k):

Your question is Vector Search from Scratch. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

Problem