📘

Top-K Retrieval

EasyRetrieval & Similarity 📚 LLM 3: Applications & Evaluation

Implement top-K retrieval using cosine similarity.

Given a query vector and a set of document vectors, find the K most similar documents.

Input:

Line 1: N D K (num documents, dimension, K)
Line 2: query vector (D floats)
Next N lines: document vectors (D floats each)

Output: Indices of top-K documents (0-based), sorted by similarity descending. Break ties by lower index first.

Example:

Input:

3 2 2
1.0 0.0
1.0 0.0
0.0 1.0
0.5 0.5

Output:

0 2

Reasoning:

The query vector is $(1.0, 0.0)$ and the document vectors are $(1.0, 0.0)$ , $(0.0, 1.0)$ , and $(0.5, 0.5)$ .
We calculate the cosine similarity between the query vector and each document vector:
- $sim_0 = \frac{(1.0, 0.0) \cdot (1.0, 0.0)}{\|(1.0, 0.0)\| \cdot \|(1.0, 0.0)\|} = 1.0$ ,
- $sim_1 = \frac{(1.0, 0.0) \cdot (0.0, 1.0)}{\|(1.0, 0.0)\| \cdot \|(0.0, 1.0)\|} = 0.0$ ,
- $sim_2 = \frac{(1.0, 0.0) \cdot (0.5, 0.5)}{\|(1.0, 0.0)\| \cdot \|(0.5, 0.5)\|} = \frac{0.5}{\sqrt{0.5}} = \frac{\sqrt{2}}{2}$ .
We sort the documents by their similarity in descending order and then by their index in ascending order: $sim_0 = 1.0$ (index 0), $sim_2 = \frac{\sqrt{2}}{2}$ (index 2), $sim_1 = 0.0$ (index 1).
The top-2 documents are indices 0 and 2, so the output is: 0 2

Constraints:

1 <= K <= N <= 100
1 <= D <= 50
Cosine similarity: dot(a,b) / (norm(a) * norm(b))

More from LLM 3: Applications & Evaluation

Beam Search DecoderHard INT8 Quantization SimulatorMedium Nucleus (Top-P) SamplingMedium Top-K SamplingEasy Temperature SamplingEasy

Background Knowledge

The problem of Top-K Retrieval is a fundamental concept in information retrieval and similarity search. It involves finding the most similar items (in this case, documents) to a given query, based on their vector representations. The cosine similarity measure is commonly used for this purpose, which calculates the cosine of the angle between two vectors. The cosine similarity between two vectors $\mathbf{a}$ and $\mathbf{b}$ is given by $\frac{\mathbf{a} \cdot \mathbf{b}}{\|\mathbf{a}\| \|\mathbf{b}\|}$ , where $\cdot$ denotes the dot product and $\| \|$ denotes the Euclidean norm.

In the context of document retrieval, each document is typically represented as a vector in a high-dimensional space, where each dimension corresponds to a feature (e.g., word frequency). The query vector is also represented in the same space. By calculating the cosine similarity between the query vector and each document vector, we can determine the similarity between the query and each document. The goal of top-K retrieval is to find the K most similar documents, ranked in descending order of similarity.

The vector space model is a fundamental concept in information retrieval, which represents documents and queries as vectors in a high-dimensional space. This allows us to apply various mathematical operations and similarity measures, such as cosine similarity, to compare and rank documents. Understanding the vector space model and cosine similarity is essential for solving the top-K retrieval problem.

Algorithm/Approach

The general approach to solving the top-K retrieval problem involves the following steps:

Calculate the cosine similarity between the query vector and each document vector
Rank the documents in descending order of similarity
Select the top-K documents with the highest similarity

This approach can be implemented using various algorithms and data structures, such as arrays, lists, or heaps. The choice of algorithm and data structure depends on the specific requirements of the problem, such as the size of the document collection and the value of K.

Step-by-Step Strategy

To implement the solution, follow these steps:

Read the input parameters: N (number of documents), D (dimension), and K (number of top documents to retrieve)
Read the query vector and store it in a suitable data structure (e.g., array or list)
Read each document vector and calculate the cosine similarity with the query vector
Store the similarity values and corresponding document indices in a suitable data structure (e.g., array or list of pairs)
Sort the documents in descending order of similarity, breaking ties by lower index first
Select the top-K documents with the highest similarity and output their indices

Common Pitfalls

When implementing the solution, watch out for the following common pitfalls:

Incorrect calculation of cosine similarity, such as forgetting to normalize the vectors or using the wrong formula
Inefficient sorting or ranking algorithms, such as using a simple bubble sort or insertion sort for large datasets
Failure to break ties correctly, such as using the wrong comparison function or not considering the document index

Time & Space Complexity

The time complexity of the solution depends on the specific algorithm and data structure used. In general, the time complexity can be broken down into the following components:

Calculating cosine similarity: O(N * D), where N is the number of documents and D is the dimension
Sorting or ranking documents: O(N log N) or O(N * K), depending on the algorithm used
Selecting top-K documents: O(K)

The space complexity depends on the data structure used to store the document vectors, similarity values, and indices. In general, the space complexity can be O(N * D) for storing the document vectors and O(N) for storing the similarity values and indices.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.