Efficient Top-K Vector Search | Dataford Interview Questions

DataArt's AI solutions may need to retrieve the most relevant embedding vectors for a query. Implement an efficient exact top-k search using cosine similarity and a bounded min-heap.

Given a query vector, a list of candidate vectors, and an integer k, return the indices of the k candidates with the highest cosine similarity to the query. Return indices in descending similarity order. If similarities are equal, return the smaller index first.

Formal Specification

Implement top_k_similar(vectors, query, k), where vectors is a list of equal-length numeric vectors, query is a numeric vector with the same dimension, and k is a positive integer. Return a list of at most k integer indices. A zero vector has cosine similarity 0 with every vector, including another zero vector.

Constraints

1 <= len(vectors) <= 200000
1 <= len(query) <= 100
Every vector has the same dimension as query
1 <= k <= len(vectors)
Vector components are finite real numbers in [-10^4, 10^4]
Zero vectors have similarity 0

Function Signature

def top_k_similar(vectors, query, k):

Your question is Efficient Top-K Vector Search. 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