Embedding Nearest Neighbors

mediumFREEWord EmbeddingsNLP 1: Foundations

Find the k nearest neighbors of a query word in an embedding space using cosine similarity.

Implement nearest_neighbors(query, k, embeddings):

query: the query word (a string that is a key of embeddings)
k: the number of neighbors to return
embeddings: a dict mapping each word to its embedding vector (a list of floats; all vectors have the same length)

Output: Return a list of the k nearest words (excluding the query), sorted by descending cosine similarity to the query's vector.

Example

Input

cat 2
4
cat 0.9 0.1
dog 0.8 0.2
car 0.1 0.9
kitten 0.85 0.15

Output

['kitten', 'dog']

Reasoning

The inputs are a query word (cat), a count k = 2, and a vocabulary where every word owns a 2-dimensional vector. cat is $(0.9, 0.1)$ .
Cosine similarity compares directions only: $\cos(q, v) = \frac{q \cdot v}{\lVert q \rVert \, \lVert v \rVert}$ . Here $\lVert \text{cat} \rVert = \sqrt{0.82} \approx 0.9055$ .
dog $(0.8, 0.2)$ : $q \cdot v = 0.72 + 0.02 = 0.74$ , $\lVert v \rVert = \sqrt{0.68} \approx 0.8246$ , so $\cos \approx 0.74 / 0.7467 \approx 0.9910$ .
kitten $(0.85, 0.15)$ : $q \cdot v = 0.765 + 0.015 = 0.78$ , $\lVert v \rVert = \sqrt{0.745} \approx 0.8631$ , so $\cos \approx 0.78 / 0.7816 \approx 0.9980$ .
car $(0.1, 0.9)$ : $q \cdot v = 0.18$ and $\lVert v \rVert = 0.9055$ , so $\cos = 0.18 / 0.82 \approx 0.2195$ .
cat itself would score exactly 1, but the query is never its own neighbour. Ranking the rest gives kitten, dog, car; the top 2 yields ['kitten', 'dog'].
Notice that kitten and dog differ by under 0.007: cosine ignores vector length, so only the angle decides.

Constraints

Exclude the query word itself from results
Sort by cosine similarity (highest first)
k is always <= number of non-query words
Use cosine similarity

Learn this in

Tokenization & EmbeddingsLLM Study Plan Retrieval-Augmented GenerationLLM Study Plan NLP FundamentalsML Study Plan Introduction to NLPNLP Study Plan

More from NLP 1: Foundations

Embedding Lookupeasy Cosine Similarity Calculatoreasy Binary Vectorizereasy Document Term Matrixeasy Character Frequency Countereasy

Concept primer

Background Knowledge

The problem involves word embeddings, which are dense vector representations of words in a high-dimensional space. These embeddings are learned such that semantically similar words are closer together in the vector space. The goal is to find the $k$ nearest neighbors of a query word based on cosine similarity, a measure of similarity between two vectors. Cosine similarity is defined as the dot product of two vectors divided by the product of their magnitudes: $\frac{\mathbf{a} \cdot \mathbf{b}}{\|\mathbf{a}\| \|\mathbf{b}\|}$ .

In the context of natural language processing (NLP), word embeddings like Word2Vec and GloVe are commonly used to capture the semantic relationships between words. These embeddings can be used for various tasks, including text classification, sentiment analysis, and information retrieval. The cosine similarity metric is particularly useful for comparing the similarity between word embeddings because it captures the directional similarity between vectors, ignoring their magnitudes.

Understanding how to calculate and utilize cosine similarity is crucial for solving this problem. Additionally, familiarity with vector operations and distance metrics will be helpful. The problem requires processing a list of words with their corresponding embedding vectors, computing similarities, and sorting the results. This involves basic programming concepts, such as data structures (e.g., lists, dictionaries) and algorithms (e.g., sorting).

Algorithm/Approach

The general approach to solving this problem involves the following algorithm pattern:

Read and parse the input data, including the query word, the value of $k$ , and the list of words with their embeddings.
Calculate the cosine similarity between the query word's embedding and each of the other word embeddings.
Sort the words based on their cosine similarity to the query word in descending order.
Select and return the top $k$ words, excluding the query word itself.

This approach leverages the concept of cosine similarity as a measure of proximity in the vector space and utilizes basic sorting algorithms to find the nearest neighbors.

Step-by-Step Strategy

To implement the solution:

Read Input: Read the query word and $k$ from the first line, and the number of words $N$ from the second line.
Parse Embeddings: Iterate through the remaining lines to parse each word and its embedding vector.
Compute Cosine Similarity: For each word, calculate the cosine similarity between its embedding and the query word's embedding.
Store Similarities: Store each word and its corresponding cosine similarity in a data structure (e.g., list of tuples).
Sort Similarities: Sort the list of words based on their cosine similarities in descending order.
Select Nearest Neighbors: Select the top $k$ words, excluding the query word, and return them as a list.

Common Pitfalls

When implementing the solution, watch out for:

Incorrectly calculating the cosine similarity, especially when dealing with vector operations.
Failing to exclude the query word from the list of nearest neighbors.
Not handling edge cases, such as when $k$ is larger than the number of available words.
Inefficient sorting or searching algorithms that could lead to performance issues for large inputs.

Time & Space Complexity

The expected time complexity for this problem is $O(N \log N)$ due to the sorting operation, where $N$ is the number of words. The space complexity is $O(N)$ , as we need to store the embeddings and similarities for all words. However, these complexities can vary depending on the specific implementation details, such as the choice of data structures and algorithms for sorting and searching.

solution.py

Python 3.13

Test results