Vantage-point tree in Python
Binary metric tree that splits points by distance from a chosen vantage point.
- File
vantage_point_tree.py, 89 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/vantage-point-tree/. The authors are listed in prov/ET45HH.jsonl.
import math
import random
from dataclasses import dataclass
from typing import Callable, Generic, Optional, Sequence, TypeVar
T = TypeVar("T")
Metric = Callable[[T, T], float]
@dataclass
class Node(Generic[T]):
vp: T
mu: float
inside: Optional["Node[T]"]
outside: Optional["Node[T]"]
class VPTree(Generic[T]):
"""A vantage-point tree over a finite set of objects of a metric space.
choose(S) picks the vantage point of a nonempty list S; the default is a random element.
The counter evaluations counts calls of the metric.
"""
def __init__(self, points: Sequence[T], d: Metric, choose: Optional[Callable[[list[T]], int]] = None, seed: int = 0) -> None:
self._d = d
self.evaluations = 0
rng = random.Random(seed)
self._choose = choose or (lambda s: rng.randrange(len(s)))
self.root = self._build(list(points))
def dist(self, a: T, b: T) -> float:
self.evaluations += 1
return self._d(a, b)
def _build(self, s: list[T]) -> Optional[Node[T]]:
if not s: # @line L1
return None
i = self._choose(s) # @line L2
vp = s[i]
rest = s[:i] + s[i + 1 :]
keyed = sorted(((self.dist(vp, x), j) for j, x in enumerate(rest))) # @line L3
m = (len(keyed) + 1) // 2
if not keyed: # @line L4
return Node(vp, 0.0, None, None)
mu = keyed[m - 1][0] # @line L5
inside = [rest[j] for _, j in keyed[:m]]
outside = [rest[j] for _, j in keyed[m:]]
return Node(vp, mu, self._build(inside), self._build(outside)) # @line L6
def range(self, q: T, r: float) -> list[T]:
"""All stored objects x with d(q, x) <= r."""
out: list[T] = []
self._range(self.root, q, r, out)
return out
def _range(self, node: Optional[Node[T]], q: T, r: float, out: list[T]) -> None:
if node is None: # @line L7
return
t = self.dist(q, node.vp) # @line L8
if t <= r:
out.append(node.vp)
if t - r <= node.mu: # @line L9
self._range(node.inside, q, r, out)
if t + r >= node.mu: # @line L10
self._range(node.outside, q, r, out)
def nearest(self, q: T) -> Optional[tuple[T, float]]:
"""A stored object at the least distance from q, with that distance; None if empty."""
best: list = [None, math.inf]
self._nearest(self.root, q, best)
return None if self.root is None else (best[0], best[1])
def _nearest(self, node: Optional[Node[T]], q: T, best: list) -> None:
if node is None: # @line L11
return
t = self.dist(q, node.vp) # @line L12
if t < best[1]:
best[0], best[1] = node.vp, t
if t <= node.mu: # @line L13
order = [(node.inside, True), (node.outside, False)]
else:
order = [(node.outside, False), (node.inside, True)]
for child, is_inside in order: # @line L14
if (t - best[1] <= node.mu) if is_inside else (t + best[1] >= node.mu): # @line L15
self._nearest(child, q, best)
test_vantage_point_tree.py 79 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/vantage-point-tree/. The authors are listed in prov/ET45HH.jsonl.
import math
import random
from vantage_point_tree import VPTree
P = [(1, 1), (2, 6), (4, 4), (5, 1), (6, 7), (8, 3), (9, 8)]
def height(node) -> int:
return 0 if node is None else 1 + max(height(node.inside), height(node.outside))
def edit(a: str, b: str) -> int:
prev = list(range(len(b) + 1))
for i, ca in enumerate(a, 1):
cur = [i]
for j, cb in enumerate(b, 1):
cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
prev = cur
return prev[-1]
def test_example() -> None:
t = VPTree(P, math.dist, choose=lambda s: 0)
assert t.evaluations == 10
assert t.root.vp == (1, 1) and abs(t.root.mu - math.sqrt(26)) < 1e-12
assert t.root.inside.vp == (5, 1) and t.root.outside.vp == (8, 3)
assert height(t.root) == 3
t.evaluations = 0
assert sorted(t.range((8, 7), 2)) == [(6, 7), (9, 8)]
assert t.evaluations == 4
t.evaluations = 0
assert t.nearest((7, 6)) == ((6, 7), math.sqrt(2))
assert t.evaluations == 4
def test_edges() -> None:
empty = VPTree([], math.dist)
assert empty.range((0, 0), 5) == [] and empty.nearest((0, 0)) is None
one = VPTree([(3, 4)], math.dist)
assert one.range((0, 0), 5) == [(3, 4)] and one.range((0, 0), 4.9) == []
same = VPTree([(1, 1)] * 6, math.dist)
assert len(same.range((1, 1), 0)) == 6
assert height(same.root) == 3
def test_height_bound() -> None:
for n in range(1, 130):
t = VPTree([(i, 0) for i in range(n)], math.dist, seed=n)
assert height(t.root) == n.bit_length()
assert t.evaluations <= n * n.bit_length()
def test_random_against_scan() -> None:
rng = random.Random(1993)
for trial in range(500):
if trial % 2:
n = rng.randint(0, 40)
dim = rng.choice([1, 2, 3])
pts = [tuple(rng.randint(0, 9) for _ in range(dim)) for _ in range(n)]
d = math.dist
q = tuple(rng.uniform(-1, 10) for _ in range(dim))
r = rng.uniform(0, 6)
else:
n = rng.randint(0, 30)
pts = ["".join(rng.choice("abc") for _ in range(rng.randint(0, 5))) for _ in range(n)]
d = edit
q = "".join(rng.choice("abc") for _ in range(rng.randint(0, 5)))
r = rng.randint(0, 3)
t = VPTree(pts, d, seed=trial)
assert sorted(t.range(q, r)) == sorted(x for x in pts if d(q, x) <= r)
nn = t.nearest(q)
if pts:
assert nn is not None and nn[1] == min(d(q, x) for x in pts) and d(q, nn[0]) == nn[1]
else:
assert nn is None
What the next level needs
- the tests did not run in this build