k-nearest neighbor graph in Python
A graph with an edge from each point to its k nearest neighbors, used as an index and in clustering.
- File
k_nearest_neighbor_graph.py, 75 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/k-nearest-neighbor-graph/. The authors are listed in prov/CST612.jsonl.
"""A k-nearest neighbor graph, built by brute force, with greedy beam search."""
import heapq
import math
from typing import Callable, Sequence
Point = Sequence[float]
def euclidean(a: Point, b: Point) -> float:
return math.dist(a, b)
class KNNGraph:
"""Directed graph with an edge from each point to its k nearest other points.
Ties in distance are broken by the smaller index.
"""
def __init__(self, points: Sequence[Point], k: int,
dist: Callable[[Point, Point], float] = euclidean) -> None:
self.points = list(points)
self.k = k
self.dist = dist
self.adj: list[list[int]] = []
n = len(self.points)
for i in range(n): # @line L1
heap: list[tuple[float, int]] = [] # max-heap of (-d, -j)
for j in range(n): # @line L2
if j == i:
continue
d = dist(self.points[i], self.points[j])
item = (-d, -j)
if len(heap) < k: # @line L3
heapq.heappush(heap, item)
elif item > heap[0]:
heapq.heapreplace(heap, item)
self.adj.append([-mj for _, mj in sorted(heap, reverse=True)]) # @line L4
def neighbors(self, i: int) -> list[int]:
return self.adj[i]
def reverse(self) -> list[list[int]]:
"""In-neighbors of each point: the reverse k-nearest neighbor lists."""
rev: list[list[int]] = [[] for _ in self.points]
for i, out in enumerate(self.adj):
for j in out:
rev[j].append(i)
return rev
def search(self, q: Point, k: int, ef: int, starts: Sequence[int]) -> list[tuple[float, int]]:
"""Beam search from the start points; returns up to k (distance, index) pairs."""
d = self.dist
visited = set(starts) # @line L5
cand = [(d(q, self.points[s]), s) for s in visited]
heapq.heapify(cand)
best = [(-dist_s, s) for dist_s, s in heapq.nsmallest(ef, cand)]
heapq.heapify(best)
while cand: # @line L6
dc, c = heapq.heappop(cand) # @line L7
if dc > -best[0][0]:
break
for e in self.adj[c]: # @line L8
if e in visited:
continue
visited.add(e)
de = d(q, self.points[e])
if len(best) < ef or de < -best[0][0]:
heapq.heappush(cand, (de, e))
heapq.heappush(best, (-de, e))
if len(best) > ef:
heapq.heappop(best)
return sorted((-nd, e) for nd, e in best)[:k] # @line L9
test_k_nearest_neighbor_graph.py 69 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/k-nearest-neighbor-graph/. The authors are listed in prov/CST612.jsonl.
import math
import random
from k_nearest_neighbor_graph import KNNGraph
LINE = [(0, 0), (1, 0), (3, 0), (6, 0), (10, 0), (11, 0)]
def oracle_lists(points, k):
out = []
for i, p in enumerate(points):
others = sorted((math.dist(p, x), j) for j, x in enumerate(points) if j != i)
out.append([j for _, j in others[:k]])
return out
def reachable(adj, starts):
seen = set(starts)
stack = list(starts)
while stack:
u = stack.pop()
for v in adj[u]:
if v not in seen:
seen.add(v)
stack.append(v)
return seen
def test_example() -> None:
g = KNNGraph(LINE, 2)
assert g.adj == [[1, 2], [0, 2], [1, 0], [2, 4], [5, 3], [4, 3]]
assert [sorted(r) for r in g.reverse()] == [[1, 2], [0, 2], [0, 1, 3], [4, 5], [3, 5], [4]]
assert g.search((9, 0), 1, ef=2, starts=[0]) == [(6.0, 2)]
assert g.search((9, 0), 1, ef=2, starts=[5]) == [(1.0, 4)]
assert reachable(g.adj, [0]) == {0, 1, 2}
def test_edge_cases() -> None:
assert KNNGraph([], 3).adj == []
assert KNNGraph([(0, 0)], 3).adj == [[]]
assert KNNGraph([(0, 0), (1, 1)], 5).adj == [[1], [0]]
def test_random_against_oracle() -> None:
rng = random.Random(612)
for _ in range(500):
n = rng.randint(1, 18)
dim = rng.randint(1, 3)
pts = [tuple(rng.randint(0, 6) for _ in range(dim)) for _ in range(n)]
k = rng.randint(1, 5)
g = KNNGraph(pts, k)
assert g.adj == oracle_lists(pts, k)
q = tuple(rng.uniform(-1, 7) for _ in range(dim))
starts = rng.sample(range(n), rng.randint(1, min(3, n)))
kk = rng.randint(1, 4)
# Theorem: with ef >= n the search is exact on the reachable set.
got = g.search(q, kk, ef=n, starts=starts)
reach = reachable(g.adj, starts)
want = sorted((math.dist(q, pts[j]), j) for j in reach)[:kk]
assert got == want
# Lemma: with any ef the first answer is a local minimum.
ef = rng.randint(1, 4)
res = g.search(q, kk, ef=ef, starts=starts)
d0, w = res[0]
assert all(math.dist(q, pts[e]) >= d0 for e in g.adj[w])
assert 1 <= len(res) <= min(kk, ef) and res == sorted(res)
What the next level needs
- the tests did not run in this build