Cosine similarity from scratch

easyRetrievalembeddingsretrievalrag

The distance function behind every vector search and RAG retriever. Two vectors in, one number in [-1, 1] out — and the edge cases are where interviews live.

Problem

The problem

Implement cosine_similarity(a, b) for two equal-length vectors given as lists of numbers. Return a single float in [-1, 1]:

cosine(a, b) = dot(a, b) / (‖a‖ · ‖b‖)

Handle the zero vector: if either input has zero magnitude the similarity is undefined — return 0.0 rather than dividing by zero. Pure Python, math available.

Concept

Cosine similarity measures the angle between two vectors, ignoring their magnitude: dot(a, b) / (‖a‖ · ‖b‖). It is 1 when the vectors point the same way, 0 when orthogonal, -1 when opposite. RAG and semantic search rank documents by cosine similarity to the query embedding because it is scale-invariant — a long document and a short one are compared on direction, not length. The interview edge case is the zero vector: its norm is 0, so the formula divides by zero. Real code returns 0 (or guards) rather than crashing.

Hints

cosine(a, b) = dot(a, b) / (norm(a) * norm(b)), where norm(v) = sqrt(sum of v_i squared).
Guard the zero vector — if either norm is 0 the similarity is undefined; return 0.0 rather than dividing by zero.

Solution

Try it yourself first — the reveal takes two clicks.

import math


def cosine_similarity(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)

solution.pyPython · Pyodide

No browser Python? Run it locally instead.

The runner is a Pyodide worker (~10 MB, WebAssembly) loaded on first run. If your network or browser blocks it, copy your code and the tests below into solution.py andtest_solution.py and run pytest.

# test_solution.py
def test_identical_vectors_are_one():
    assert abs(cosine_similarity([1, 2, 3], [1, 2, 3]) - 1.0) < 1e-9

def test_orthogonal_vectors_are_zero():
    assert abs(cosine_similarity([1, 0], [0, 1])) < 1e-9

def test_opposite_vectors_are_minus_one():
    assert abs(cosine_similarity([1, 0], [-1, 0]) + 1.0) < 1e-9

def test_scale_invariant():
    # b is a scaled copy of a, so direction is identical -> similarity 1.
    assert abs(cosine_similarity([1, 2, 3], [2, 4, 6]) - 1.0) < 1e-9

def test_zero_vector_is_safe():
    assert cosine_similarity([0, 0], [1, 1]) == 0.0