PIXELBANKv8.2.1
Menu

Embedding Nearest Neighbors

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

Input format:

  • Line 1: The query word and k (space-separated)
  • Line 2: Number of words N
  • Lines 3 to N+2: word followed by its embedding vector

Output: List of k nearest words (excluding the query), sorted by descending cosine similarity. Print as a Python list of strings.

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:

Compute cosine similarity of each word with "cat" [0.9, 0.1]:

  • dog [0.8, 0.2]: dot=0.74, norms: 0.906*0.825=0.747, sim=0.74/0.747=0.9908
  • car [0.1, 0.9]: dot=0.18, norms: 0.906*0.906=0.820, sim=0.18/0.820=0.2195
  • kitten [0.85, 0.15]: dot=0.78, norms: 0.906*0.863=0.782, sim=0.78/0.782=0.9975

Top 2: kitten (0.9975), dog (0.9908)

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
Editor

Test Results

0/0
Run code to see test results.