Brute-force nearest neighbor search in Python

Compute the distance from the query to every point and keep the smallest; the baseline all indexes are measured against.

File: brute_force_nearest_neighbor_search.py, 35 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

brute_force_nearest_neighbor_search.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/brute-force-nearest-neighbor-search/. The authors are listed in prov/VXFMJM.jsonl.

import heapq
from collections.abc import Callable, Sequence

Point = Sequence[float]


def squared_euclidean(a: Point, b: Point) -> float:
    """Return the squared Euclidean distance of two points of the same dimension."""
    return sum((x - y) * (x - y) for x, y in zip(a, b, strict=True))


def brute_force_knn(
    points: Sequence[Point],
    q: Point,
    k: int,
    dist: Callable[[Point, Point], float] = squared_euclidean,
) -> list[tuple[float, int]]:
    """Return the min(k, n) pairs (distance, index) of the points nearest to q, in increasing order.

    Ties in distance are broken by the smaller index. Precondition: k >= 1.
    Cost: n calls of dist and O(n log k) heap time; O(k) extra space.
    """
    if k < 1:
        raise ValueError("k must be at least 1")
    h: list[tuple[float, int]] = []  # @line L1: max-heap through negated pairs
    for i in range(len(points)):  # @line L2
        delta = dist(q, points[i])  # @line L3
        if len(h) < k:
            heapq.heappush(h, (-delta, -i))  # @line L4
        elif (delta, i) < (-h[0][0], -h[0][1]):
            heapq.heapreplace(h, (-delta, -i))  # @line L5
    return sorted((-nd, -ni) for nd, ni in h)  # @line L6

test_brute_force_nearest_neighbor_search.py 51 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/brute-force-nearest-neighbor-search/. The authors are listed in prov/VXFMJM.jsonl.

import math
import random

import pytest

from brute_force_nearest_neighbor_search import brute_force_knn, squared_euclidean

P = [(1, 1), (2, 4), (4, 2), (5, 5), (6, 1), (7, 3), (3, 6), (4, 4)]


def test_example() -> None:
    assert brute_force_knn(P, (4, 3), 3) == [(1, 2), (1, 7), (5, 1)]


def test_edges() -> None:
    assert brute_force_knn([], (0, 0), 3) == []
    assert brute_force_knn([(2, 2)], (0, 0), 1) == [(8, 0)]
    assert brute_force_knn([(1, 1)] * 4, (0, 0), 2) == [(2, 0), (2, 1)]
    assert brute_force_knn(P, (4, 3), 100) == sorted((squared_euclidean(p, (4, 3)), i) for i, p in enumerate(P))
    with pytest.raises(ValueError):
        brute_force_knn(P, (0, 0), 0)


def test_other_metric() -> None:
    def l1(a, b):
        return sum(abs(x - y) for x, y in zip(a, b))

    got = brute_force_knn(P, (4, 3), 5, l1)
    assert [i for _, i in got] == [2, 7, 1, 3, 5]


def test_decreasing_order_replaces_each_time() -> None:
    pts = [(float(10 - i),) for i in range(10)]
    assert brute_force_knn(pts, (0.0,), 2) == [(1.0, 9), (4.0, 8)]


def test_random_against_sort() -> None:
    rng = random.Random(4242)
    for _ in range(600):
        n = rng.randint(0, 40)
        d = rng.randint(1, 4)
        pts = [tuple(rng.randint(-5, 5) for _ in range(d)) for _ in range(n)]
        q = tuple(rng.randint(-5, 5) for _ in range(d))
        k = rng.randint(1, 10)
        want = sorted((squared_euclidean(p, q), i) for i, p in enumerate(pts))[:k]
        assert brute_force_knn(pts, q, k) == want
        got_l2 = brute_force_knn(pts, q, k, lambda a, b: math.dist(a, b))
        assert [i for _, i in got_l2] == [i for _, i in want]

What the next level needs

the tests did not run in this build

Read the full page about Brute-force nearest neighbor search