Ball tree in Python
Binary tree of hyperspheres that supports nearest-neighbor search in general metric spaces.
- File
ball_tree.py, 106 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/ball-tree/. The authors are listed in prov/DT55WY.jsonl.
import heapq
import math
from collections.abc import Callable, Sequence
from typing import Generic, TypeVar
P = TypeVar("P")
def euclidean(x: Sequence[float], y: Sequence[float]) -> float:
return math.sqrt(sum((a - b) * (a - b) for a, b in zip(x, y)))
class Node(Generic[P]):
__slots__ = ("center", "radius", "points", "left", "right")
def __init__(self, center: int, radius: float) -> None:
self.center = center
self.radius = radius
self.points: list[int] | None = None
self.left: Node[P] | None = None
self.right: Node[P] | None = None
class BallTree(Generic[P]):
"""Metric ball tree. Each node has a center point and a radius that covers its subtree.
Build: O(n log n) distance evaluations, height ceil(lg(n/b)) + 1. Queries: O(n) worst case.
"""
def __init__(
self,
points: Sequence[P],
dist: Callable[[P, P], float] = euclidean,
leaf_size: int = 2,
) -> None:
if leaf_size < 1:
raise ValueError("leaf_size must be at least 1")
self.points = list(points)
self.dist = dist
self.b = leaf_size
self.root = self._build(list(range(len(self.points)))) if self.points else None
def _build(self, ids: list[int]) -> Node[P]:
d = self.dist
pts = self.points
c = ids[0] # @line L2
node = Node(c, max(d(pts[c], pts[i]) for i in ids))
if len(ids) <= self.b: # @line L1
node.points = ids
return node
a = max(ids, key=lambda i: d(pts[c], pts[i])) # @line L3
z = max(ids, key=lambda i: d(pts[a], pts[i]))
ids = sorted(ids, key=lambda i: (d(pts[i], pts[a]) - d(pts[i], pts[z]), i)) # @line L4
h = len(ids) // 2
node.left = self._build(ids[:h]) # @line L5
node.right = self._build(ids[h:])
return node
def knn(self, q: P, k: int) -> list[tuple[float, int]]:
"""Return the k pairs (distance, index) nearest to q, nearest first."""
heap: list[tuple[float, int]] = []
def tau() -> float:
return -heap[0][0] if len(heap) == k else math.inf
def visit(t: Node[P], dc: float) -> None:
if dc - t.radius >= tau(): # @line L6
return
if t.points is not None: # @line L7
for i in t.points:
di = self.dist(q, self.points[i])
if len(heap) < k:
heapq.heappush(heap, (-di, -i))
elif (-di, -i) > heap[0]:
heapq.heapreplace(heap, (-di, -i))
return
assert t.left is not None and t.right is not None
dl = self.dist(q, self.points[t.left.center]) # @line L8
dr = self.dist(q, self.points[t.right.center])
if dl <= dr:
visit(t.left, dl)
visit(t.right, dr)
else:
visit(t.right, dr)
visit(t.left, dl)
if self.root is not None and k > 0:
visit(self.root, self.dist(q, self.points[self.root.center]))
return sorted((-a, -b) for a, b in heap)
def range(self, q: P, r: float) -> list[int]:
"""Return the indices of all points within distance r of q."""
out: list[int] = []
stack = [self.root] if self.root is not None else []
while stack:
t = stack.pop()
if self.dist(q, self.points[t.center]) - t.radius > r:
continue
if t.points is not None:
out.extend(i for i in t.points if self.dist(q, self.points[i]) <= r)
else:
stack.extend((t.left, t.right))
return out
test_ball_tree.py 65 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/ball-tree/. The authors are listed in prov/DT55WY.jsonl.
import math
import random
from ball_tree import BallTree, euclidean
PTS = [(1, 1), (2, 3), (3, 1), (6, 5), (7, 7), (8, 4), (2, 2), (7, 5)]
def check(tree, t, depth=0):
ids = t.points if t.points is not None else collect(t)
assert all(tree.dist(tree.points[t.center], tree.points[i]) <= t.radius + 1e-12 for i in ids)
if t.points is None:
a, b = collect(t.left), collect(t.right)
assert abs(len(a) - len(b)) <= 1
return max(check(tree, t.left, depth + 1), check(tree, t.right, depth + 1))
assert len(t.points) <= tree.b
return depth + 1
def collect(t):
if t.points is not None:
return list(t.points)
return collect(t.left) + collect(t.right)
def test_example() -> None:
tree = BallTree(PTS, leaf_size=2)
assert tree.root.center == 0
assert math.isclose(tree.root.radius, math.sqrt(72))
assert sorted(collect(tree.root.left)) == [3, 4, 5, 7]
assert check(tree, tree.root) == 3
assert tree.knn((5, 4), 2) == [(math.sqrt(2), 3), (math.sqrt(5), 7)]
assert tree.knn((6, 6), 1) == [(1.0, 3)]
def test_random_against_brute_force() -> None:
rng = random.Random(8)
for _ in range(500):
d = rng.randint(1, 5)
n = rng.randint(1, 40)
pts = [tuple(float(rng.randint(-6, 6)) for _ in range(d)) for _ in range(n)]
tree = BallTree(pts, leaf_size=rng.randint(1, 4))
h = check(tree, tree.root)
assert h <= math.ceil(math.log2(max(1, n))) + 1
q = tuple(rng.uniform(-7, 7) for _ in range(d))
k = rng.randint(1, 6)
want = sorted(euclidean(q, p) for p in pts)[:k]
got = tree.knn(q, k)
assert [a for a, _ in got] == want
assert all(euclidean(q, pts[i]) == a for a, i in got)
r = rng.uniform(0, 8)
assert sorted(tree.range(q, r)) == [i for i, p in enumerate(pts) if euclidean(q, p) <= r]
def test_string_metric() -> None:
rng = random.Random(5)
words = ["".join(rng.choice("ab") for _ in range(8)) for _ in range(60)]
ham = lambda x, y: sum(a != b for a, b in zip(x, y))
tree = BallTree(words, dist=ham, leaf_size=3)
for _ in range(100):
q = "".join(rng.choice("ab") for _ in range(8))
assert [a for a, _ in tree.knn(q, 3)] == sorted(ham(q, w) for w in words)[:3]
What the next level needs
- the tests did not run in this build