Problem 639634 · hard · Level 06 Heuristics & Optimization

Search Only the Promising Shelves

similarity search · approximate nearest neighbours · k-means · cosine similarity · embeddings · inverted index

An online shop stores an embedding vector for every product and answers "show me products like this one" by cosine similarity. Comparing a query with every product is too slow for a big catalogue, so the shop builds an index: products are sorted onto a few "shelves" of similar products, and a query only looks at the most promising shelves.

Write approximate_search(vectors, n_lists, rounds, queries, k, probes) that returns a tuple (results, comparisons). Throughout, the similarity of two unit vectors is their dot product rounded to 9 decimal places.

Building the index.

  1. Normalise every product vector to length 1.
  2. The n_lists centres start as the unit vectors of products 0, 1, ..., n_lists - 1.
  3. Repeat rounds times: put every product on the shelf of the centre it is most similar to (the lower centre index on a tie); then replace every centre that has at least one product by the mean of its products' unit vectors, normalised to length 1 (a centre whose shelf is empty, or whose mean is all zeros, stays as it is).
  4. Finally put every product on a shelf once more, by the same rule, with the final centres.

Answering a query. Normalise the query. Rank the centres by similarity to it (highest first, lower index on a tie) and take the first probes shelves. Rank all products on those shelves by similarity to the query (highest first, lower product index on a tie) and keep the first k indices (fewer if the shelves hold fewer products).

results is the list of these index lists, one per query, and comparisons is the total number of similarities computed while answering the queries: n_lists per query for the centres, plus one per product on its probed shelves.

product_vectors(n, dim, seed) builds the catalogues and queries used by the tests; it is available in your code.

Examples

Input:  vectors = [[1, 0], [0, 1], [-1, 0.2], [0.9, 0.1], [0.1, 0.9], [-0.9, -0.1]],
        n_lists = 3, rounds = 2, queries = [[1, 0.2], [-1, 0]], k = 2, probes = 1
Output: ([[3, 0], [5, 2]], 10)
Explanation: the shelves are {0, 3}, {1, 4} and {2, 5}. Each query computes 3 centre similarities
and 2 product similarities.

Input:  the same vectors and index, queries = [[0.8, 0.6]], k = 3, probes = 2
Output: ([[3, 0, 4]], 7)
Explanation: with probes = 1 only the shelf {0, 3} would be searched and product 4 would be missed.

Constraints

  • 1 <= n_lists <= len(vectors) <= 1500, 2 to 12 dimensions, 0 <= rounds <= 6, 1 <= probes <= n_lists
  • up to 40 queries, 1 <= k <= 10; no vector is all zeros

Goals

  • Group embeddings into lists around centres found by k-means on unit vectors
  • Answer a query by scanning only the lists whose centres are closest to it
  • Count the work saved and see what approximate search can miss
Starting Python…