Interval tree search in Python
Descend an interval tree, pruning subtrees by their maximum endpoint, to report overlaps.
- File
interval_tree_search.py, 111 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/interval-tree-search/. The authors are listed in prov/C1KCSN.jsonl.
"""Search of a static augmented interval tree, for every end point type.
The tree is an implicit balanced binary search tree over the intervals sorted
by left end: the node of the index range [lo, hi) is mid = (lo + hi) // 2, its
left subtree is [lo, mid) and its right subtree is [mid + 1, hi). Each node
stores the largest right end of its subtree.
A query is given by two monotone tests:
reach(r) -- true for every right end r at or above a threshold,
start(l) -- true for every left end l at or below a threshold.
An interval matches when reach(r) and start(l). The table QUERY_TESTS gives
the tests for overlap and stabbing queries on closed, half-open and open
intervals.
"""
from typing import Callable, Optional, Sequence, Tuple
Interval = Tuple[float, float]
KINDS = ("closed", "half-open", "open")
def overlap_tests(a, b, kind="closed"):
"""Tests for the query interval <a, b> of the same type as the stored ones."""
if kind == "closed": # [l, r] meets [a, b] iff a <= r and l <= b
return (lambda r: r >= a), (lambda l: l <= b)
if kind in ("half-open", "open"): # strict at both cross conditions
return (lambda r: r > a), (lambda l: l < b)
raise ValueError("kind must be closed, half-open or open")
def stab_tests(p, kind="closed"):
"""Tests for the query point p."""
if kind == "closed": # l <= p <= r
return (lambda r: r >= p), (lambda l: l <= p)
if kind == "half-open": # l <= p < r
return (lambda r: r > p), (lambda l: l <= p)
if kind == "open": # l < p < r
return (lambda r: r > p), (lambda l: l < p)
raise ValueError("kind must be closed, half-open or open")
class IntervalTreeSearch:
"""Static interval tree; queries return indices into the input list."""
def __init__(self, intervals: Sequence[Interval]):
self.iv = list(intervals)
self.order = sorted(range(len(self.iv)), key=lambda j: (self.iv[j][0], self.iv[j][1], j))
n = len(self.order)
self.max = [0.0] * n # max right end in the subtree whose root is position m
self._build(0, n)
def _build(self, lo, hi):
if lo >= hi:
return None
mid = (lo + hi) // 2
best = self.iv[self.order[mid]][1]
for child in (self._build(lo, mid), self._build(mid + 1, hi)):
if child is not None and child > best:
best = child
self.max[mid] = best
return best
# REPORT: all matching intervals, in the order of the tree.
def report(self, reach: Callable, start: Callable) -> list[int]:
out: list[int] = []
self._report(0, len(self.order), reach, start, out)
return out
def _report(self, lo, hi, reach, start, out):
if lo >= hi: # @line L1
return
mid = (lo + hi) // 2
if not reach(self.max[mid]): # @line L1
return
self._report(lo, mid, reach, start, out) # @line L2
j = self.order[mid]
l, r = self.iv[j]
if not start(l): # @line L3
return
if reach(r): # @line L4
out.append(j)
self._report(mid + 1, hi, reach, start, out) # @line L5
# SEARCH: one matching interval, or None.
def search(self, reach: Callable, start: Callable) -> Optional[int]:
lo, hi = 0, len(self.order) # @line L6
while lo < hi: # @line L7
mid = (lo + hi) // 2
j = self.order[mid]
l, r = self.iv[j]
if reach(r) and start(l): # @line L7
return j
left_mid = (lo + mid) // 2
if lo < mid and reach(self.max[left_mid]): # @line L8
hi = mid
else:
lo = mid + 1
return None # @line L9
def overlap(self, a, b, kind="closed") -> list[int]:
return self.report(*overlap_tests(a, b, kind))
def stab(self, p, kind="closed") -> list[int]:
return self.report(*stab_tests(p, kind))
def interval_tree_search(intervals: Sequence[Interval], a, b, kind="closed") -> list[int]:
"""Indices of the intervals that meet the query <a, b>, by building the tree once."""
return IntervalTreeSearch(intervals).overlap(a, b, kind)
test_interval_tree_search.py 155 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/interval-tree-search/. The authors are listed in prov/C1KCSN.jsonl.
import random
from interval_tree_search import (
IntervalTreeSearch,
interval_tree_search,
overlap_tests,
stab_tests,
)
KINDS = ("closed", "half-open", "open")
def members(x, l, r, kind):
if kind == "closed":
return l <= x <= r
if kind == "half-open":
return l <= x < r
return l < x < r
def meets(iv, q, kind):
"""Brute force on a fine grid: integer end points, so quarter steps suffice."""
lo = min(iv[0], q[0]) - 1
hi = max(iv[1], q[1]) + 1
x = lo
while x <= hi:
if members(x, *iv, kind) and members(x, *q, kind):
return True
x += 0.25
return False
EXAMPLE = [(1, 3), (2, 5), (3, 4), (5, 8), (6, 7), (8, 9), (7, 10)]
def names(t, idx):
return sorted(t.iv[j] for j in idx)
def test_example_overlap():
t = IntervalTreeSearch(EXAMPLE)
assert names(t, t.overlap(4, 6, "closed")) == [(2, 5), (3, 4), (5, 8), (6, 7)]
assert names(t, t.overlap(4, 6, "half-open")) == [(2, 5), (5, 8)]
assert names(t, t.overlap(4, 6, "open")) == [(2, 5), (5, 8)]
# in tree order, the order of the left ends
assert t.overlap(4, 6, "closed") == [1, 2, 3, 4]
def test_example_stab():
t = IntervalTreeSearch(EXAMPLE)
assert names(t, t.stab(5, "closed")) == [(2, 5), (5, 8)]
assert names(t, t.stab(5, "half-open")) == [(5, 8)]
assert names(t, t.stab(5, "open")) == []
assert names(t, t.stab(7, "closed")) == [(5, 8), (6, 7), (7, 10)]
assert names(t, t.stab(7, "half-open")) == [(5, 8), (7, 10)]
assert names(t, t.stab(7, "open")) == [(5, 8)]
def test_example_search():
t = IntervalTreeSearch(EXAMPLE)
assert t.iv[t.search(*overlap_tests(4, 6))] == (5, 8)
assert t.iv[t.search(*overlap_tests(9, 9))] == (7, 10)
assert t.search(*overlap_tests(11, 12)) is None
assert t.search(*overlap_tests(10, 12, "half-open")) is None
assert t.iv[t.search(*overlap_tests(10, 12, "closed"))] == (7, 10)
def test_touching():
for kind, want in (("closed", [0, 1]), ("half-open", []), ("open", [])):
assert interval_tree_search([(1, 3), (5, 6)], 3, 5, kind) == want
t = IntervalTreeSearch([(1, 3), (3, 5)])
assert t.stab(3, "closed") == [0, 1]
assert t.stab(3, "half-open") == [1]
assert t.stab(3, "open") == []
def test_edge_cases():
t = IntervalTreeSearch([])
assert t.overlap(0, 1) == [] and t.stab(0) == [] and t.search(*stab_tests(0)) is None
t = IntervalTreeSearch([(2, 2)])
assert t.stab(2) == [0] and t.stab(1) == [] and t.overlap(2, 2) == [0]
# duplicates and a query inside one interval
t = IntervalTreeSearch([(0, 10), (0, 10), (4, 5)])
assert sorted(t.overlap(6, 7)) == [0, 1]
assert sorted(t.stab(4, "half-open")) == [0, 1, 2]
def brute_overlap(ivs, a, b, kind):
return sorted(j for j, iv in enumerate(ivs) if meets(iv, (a, b), kind))
def brute_stab(ivs, p, kind):
return sorted(j for j, (l, r) in enumerate(ivs) if members(p, l, r, kind))
def random_family(rng, kind):
n = rng.randint(0, 25)
out = []
for _ in range(n):
l = rng.randint(0, 20)
r = l + rng.randint(0 if kind == "closed" else 1, 8)
out.append((l, r))
return out
def test_random_against_brute_force():
rng = random.Random(20261105)
for _ in range(600):
kind = rng.choice(KINDS)
ivs = random_family(rng, kind)
t = IntervalTreeSearch(ivs)
a = rng.randint(-1, 29)
b = a + rng.randint(0 if kind == "closed" else 1, 6)
want = brute_overlap(ivs, a, b, kind)
got = t.overlap(a, b, kind)
assert sorted(got) == want and len(set(got)) == len(got)
one = t.search(*overlap_tests(a, b, kind))
assert (one is None) == (not want) and (one is None or one in want)
p = rng.randint(-1, 29) + rng.choice((0, 0, 0.5))
assert sorted(t.stab(p, kind)) == brute_stab(ivs, p, kind)
one = t.search(*stab_tests(p, kind))
assert (one is None) == (not brute_stab(ivs, p, kind))
def test_output_sensitive_visits():
"""REPORT never expands more than (k + 1) * (h + 1) nodes."""
rng = random.Random(7)
for _ in range(100):
ivs = random_family(rng, "closed")
t = IntervalTreeSearch(ivs)
n = len(ivs)
h = n.bit_length()
a = rng.randint(0, 28)
b = a + rng.randint(0, 4)
reach, start = overlap_tests(a, b)
expanded = 0
# count nodes that pass L1 by replaying the recursion
def walk(lo, hi):
nonlocal expanded
if lo >= hi:
return
mid = (lo + hi) // 2
if not reach(t.max[mid]):
return
expanded += 1
walk(lo, mid)
if start(t.iv[t.order[mid]][0]):
walk(mid + 1, hi)
walk(0, n)
k = len(t.overlap(a, b))
assert expanded <= (k + 1) * max(h, 1)
What the next level needs
- the tests did not run in this build