Finger tree in Python
A persistent balanced sequence with amortized O(1) access at both ends and O(log n) split and concatenation.
- File
finger_tree.py, 330 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/finger-tree/. The authors are listed in prov/BJ50ZK.jsonl.
"""A 2-3 finger tree (Hinze and Paterson) measured by size.
This version is strict: the middle tree is built at once, not suspended.
Each operation returns a new tree and never changes an old one.
"""
from collections.abc import Iterator
from typing import Any, Generic, Optional, TypeVar, Union
T = TypeVar("T")
class _Node:
"""A 2-3 node whose children are items (level 0) or nodes of one level less."""
__slots__ = ("size", "items")
def __init__(self, items: tuple[Any, ...]) -> None:
self.items = items
self.size = sum(_sz(x) for x in items)
def _sz(x: Any) -> int:
return x.size if isinstance(x, _Node) else 1
def _dsz(xs: tuple[Any, ...]) -> int:
return sum(_sz(x) for x in xs)
class _Empty:
__slots__ = ()
size = 0
class _Single:
__slots__ = ("x", "size")
def __init__(self, x: Any) -> None:
self.x = x
self.size = _sz(x)
class _Deep:
__slots__ = ("size", "pr", "m", "sf")
def __init__(self, pr: tuple[Any, ...], m: "_Tree", sf: tuple[Any, ...]) -> None:
self.pr, self.m, self.sf = pr, m, sf
self.size = _dsz(pr) + m.size + _dsz(sf)
_Tree = Union[_Empty, _Single, _Deep]
EMPTY = _Empty()
def _cons(a: Any, t: _Tree) -> _Tree:
if isinstance(t, _Empty): # @line L1
return _Single(a)
if isinstance(t, _Single): # @line L2
return _Deep((a,), EMPTY, (t.x,))
pr, m, sf = t.pr, t.m, t.sf # @line L3
if len(pr) == 4: # @line L4
b, c, d, e = pr
return _Deep((a, b), _cons(_Node((c, d, e)), m), sf)
return _Deep((a,) + pr, m, sf) # @line L5
def _snoc(t: _Tree, a: Any) -> _Tree:
if isinstance(t, _Empty):
return _Single(a)
if isinstance(t, _Single):
return _Deep((t.x,), EMPTY, (a,))
pr, m, sf = t.pr, t.m, t.sf
if len(sf) == 4:
e, d, c, b = sf
return _Deep(pr, _snoc(m, _Node((e, d, c))), (b, a))
return _Deep(pr, m, sf + (a,))
def _cons_all(xs: tuple[Any, ...], t: _Tree) -> _Tree:
for x in reversed(xs): # @line L6
t = _cons(x, t)
return t
def _snoc_all(t: _Tree, xs: tuple[Any, ...]) -> _Tree:
for x in xs: # @line L7
t = _snoc(t, x)
return t
def _tree_of(xs: tuple[Any, ...]) -> _Tree:
return _cons_all(xs, EMPTY) # @line L8
def _viewl(t: _Tree) -> Optional[tuple[Any, _Tree]]:
if isinstance(t, _Empty): # @line L9
return None
if isinstance(t, _Single): # @line L10
return (t.x, EMPTY)
a = t.pr[0] # @line L11
return (a, _deepl(t.pr[1:], t.m, t.sf)) # @line L12
def _viewr(t: _Tree) -> Optional[tuple[_Tree, Any]]:
if isinstance(t, _Empty):
return None
if isinstance(t, _Single):
return (EMPTY, t.x)
a = t.sf[-1]
return (_deepr(t.pr, t.m, t.sf[:-1]), a)
def _deepl(pr: tuple[Any, ...], m: _Tree, sf: tuple[Any, ...]) -> _Tree:
if pr: # @line L13
return _Deep(pr, m, sf)
v = _viewl(m) # @line L14
if v is None:
return _tree_of(sf)
node, m2 = v # @line L15
return _Deep(node.items, m2, sf)
def _deepr(pr: tuple[Any, ...], m: _Tree, sf: tuple[Any, ...]) -> _Tree:
if sf:
return _Deep(pr, m, sf)
v = _viewr(m)
if v is None:
return _tree_of(pr)
m2, node = v
return _Deep(pr, m2, node.items)
def _nodes(xs: tuple[Any, ...]) -> tuple[_Node, ...]:
n = len(xs)
if n == 2: # @line L16
return (_Node(xs),)
if n == 3: # @line L17
return (_Node(xs),)
if n == 4: # @line L18
return (_Node(xs[:2]), _Node(xs[2:]))
return (_Node(xs[:3]),) + _nodes(xs[3:]) # @line L19
def _app3(t1: _Tree, ts: tuple[Any, ...], t2: _Tree) -> _Tree:
if isinstance(t1, _Empty): # @line L20
return _cons_all(ts, t2)
if isinstance(t2, _Empty): # @line L21
return _snoc_all(t1, ts)
if isinstance(t1, _Single): # @line L22
return _cons(t1.x, _cons_all(ts, t2))
if isinstance(t2, _Single): # @line L23
return _snoc(_snoc_all(t1, ts), t2.x)
mid = _app3(t1.m, _nodes(t1.sf + ts + t2.pr), t2.m) # @line L24
return _Deep(t1.pr, mid, t2.sf)
def _locate(i: int, xs: tuple[Any, ...]) -> tuple[int, int]:
"""For 0 <= i < size(xs): (k, j) with position i at offset j inside xs[k]."""
for k, x in enumerate(xs): # @line L25
s = _sz(x)
if i < s:
return (k, i)
i -= s
raise IndexError("index outside the digit")
def _split_digit(i: int, xs: tuple[Any, ...]) -> tuple[tuple[Any, ...], Any, tuple[Any, ...]]:
k, _ = _locate(i, xs)
return (xs[:k], xs[k], xs[k + 1 :])
def _split_tree(i: int, t: _Tree) -> tuple[_Tree, Any, _Tree]:
"""For 0 <= i < size(t): (l, x, r) with x the element that holds position i."""
if isinstance(t, _Single): # @line L26
return (EMPTY, t.x, EMPTY)
assert isinstance(t, _Deep)
spr = _dsz(t.pr)
if i < spr: # @line L27
l, x, r = _split_digit(i, t.pr)
return (_tree_of(l), x, _deepl(r, t.m, t.sf))
i -= spr # @line L28
if i < t.m.size:
ml, xs, mr = _split_tree(i, t.m)
l, x, r = _split_digit(i - ml.size, xs.items)
return (_deepr(t.pr, ml, l), x, _deepl(r, mr, t.sf))
l, x, r = _split_digit(i - t.m.size, t.sf) # @line L29
return (_deepr(t.pr, t.m, l), x, _tree_of(r))
def _lookup(i: int, t: _Tree) -> Any:
"""For 0 <= i < size(t): the item at position i."""
while isinstance(t, _Deep) and _dsz(t.pr) <= i < _dsz(t.pr) + t.m.size: # @line L30
i, t = i - _dsz(t.pr), t.m # @line L31
if isinstance(t, _Single): # @line L32
xs: tuple[Any, ...] = (t.x,)
else:
assert isinstance(t, _Deep)
xs = t.pr + t.sf # @line L33
if i >= _dsz(t.pr):
i -= t.m.size
while True:
k, i = _locate(i, xs) # @line L34
if not isinstance(xs[k], _Node):
return xs[k]
xs = xs[k].items
def _iter(t: _Tree) -> Iterator[Any]:
def elem(x: Any) -> Iterator[Any]:
if isinstance(x, _Node):
for c in x.items:
yield from elem(c)
else:
yield x
if isinstance(t, _Single):
yield from elem(t.x)
elif isinstance(t, _Deep):
for x in t.pr:
yield from elem(x)
for x in _iter(t.m):
yield from elem(x)
for x in t.sf:
yield from elem(x)
def _check_elem(x: Any, level: int) -> int:
if level == 0:
assert not isinstance(x, _Node), "an item expected at level 0"
return 1
assert isinstance(x, _Node), "a node expected above level 0"
assert len(x.items) in (2, 3), "a node has 2 or 3 children"
s = sum(_check_elem(c, level - 1) for c in x.items)
assert s == x.size, "cached node size"
return s
def _check_tree(t: _Tree, level: int) -> int:
if isinstance(t, _Empty):
return 0
if isinstance(t, _Single):
s = _check_elem(t.x, level)
else:
assert 1 <= len(t.pr) <= 4 and 1 <= len(t.sf) <= 4, "digits have 1 to 4 elements"
s = sum(_check_elem(x, level) for x in t.pr)
s += _check_tree(t.m, level + 1)
s += sum(_check_elem(x, level) for x in t.sf)
assert s == t.size, "cached tree size"
return s
class FingerTree(Generic[T]):
"""A persistent sequence. Every method returns new trees; self never changes."""
__slots__ = ("_t",)
def __init__(self, _t: _Tree = EMPTY) -> None:
self._t = _t
def __len__(self) -> int:
return self._t.size
def __iter__(self) -> Iterator[T]:
return _iter(self._t)
def to_list(self) -> list[T]:
return list(_iter(self._t))
@classmethod
def from_iterable(cls, xs: "list[T]") -> "FingerTree[T]":
t: _Tree = EMPTY
for x in xs:
t = _snoc(t, x)
return cls(t)
def cons(self, x: T) -> "FingerTree[T]":
return FingerTree(_cons(x, self._t))
def snoc(self, x: T) -> "FingerTree[T]":
return FingerTree(_snoc(self._t, x))
def viewl(self) -> "Optional[tuple[T, FingerTree[T]]]":
v = _viewl(self._t)
return None if v is None else (v[0], FingerTree(v[1]))
def viewr(self) -> "Optional[tuple[FingerTree[T], T]]":
v = _viewr(self._t)
return None if v is None else (FingerTree(v[0]), v[1])
def head(self) -> T:
v = _viewl(self._t)
if v is None:
raise IndexError("head of an empty finger tree")
return v[0]
def last(self) -> T:
v = _viewr(self._t)
if v is None:
raise IndexError("last of an empty finger tree")
return v[1]
def append(self, other: "FingerTree[T]") -> "FingerTree[T]":
return FingerTree(_app3(self._t, (), other._t))
def split_at(self, i: int) -> "tuple[FingerTree[T], FingerTree[T]]":
"""The first i items and the rest."""
if i <= 0:
return (FingerTree(EMPTY), self)
if i >= self._t.size:
return (self, FingerTree(EMPTY))
l, x, r = _split_tree(i, self._t)
return (FingerTree(l), FingerTree(_cons(x, r)))
def lookup(self, i: int) -> T:
if not 0 <= i < self._t.size:
raise IndexError("finger tree index out of range")
return _lookup(i, self._t)
def depth(self) -> int:
d, t = 0, self._t
while isinstance(t, _Deep):
d, t = d + 1, t.m
return d
def check_invariant(self) -> None:
_check_tree(self._t, 0)
test_finger_tree.py 139 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/finger-tree/. The authors are listed in prov/BJ50ZK.jsonl.
import random
import pytest
from finger_tree import FingerTree
def test_example() -> None:
t: FingerTree[int] = FingerTree()
shapes = []
for x in range(1, 10):
t = t.snoc(x)
t.check_invariant()
shapes.append(t.depth())
assert t.to_list() == list(range(1, 10))
assert shapes == [0, 1, 1, 1, 1, 1, 1, 1, 2]
assert t.lookup(4) == 5
left, right = t.split_at(4)
assert left.to_list() == [1, 2, 3, 4]
assert right.to_list() == [5, 6, 7, 8, 9]
both = right.append(left)
both.check_invariant()
assert both.to_list() == [5, 6, 7, 8, 9, 1, 2, 3, 4]
assert t.to_list() == list(range(1, 10))
def test_empty() -> None:
t: FingerTree[int] = FingerTree()
t.check_invariant()
assert len(t) == 0 and t.viewl() is None and t.viewr() is None
with pytest.raises(IndexError):
t.head()
with pytest.raises(IndexError):
t.last()
with pytest.raises(IndexError):
t.lookup(0)
a, b = t.split_at(0)
assert a.to_list() == [] and b.to_list() == []
def test_ends() -> None:
t = FingerTree.from_iterable(list(range(100)))
out = []
while True:
v = t.viewl()
if v is None:
break
x, t = v
t.check_invariant()
out.append(x)
assert out == list(range(100))
t = FingerTree.from_iterable(list(range(100)))
out = []
while len(t):
t, x = t.viewr() # type: ignore[misc]
t.check_invariant()
out.append(x)
assert out == list(range(99, -1, -1))
def test_depth_is_logarithmic() -> None:
t: FingerTree[int] = FingerTree()
for x in range(5000):
t = t.cons(x)
t.check_invariant()
assert 2 ** (t.depth() - 1) <= len(t)
def test_all_splits_and_appends_small() -> None:
for n in range(0, 40):
t = FingerTree.from_iterable(list(range(n)))
for i in range(-1, n + 2):
a, b = t.split_at(i)
a.check_invariant()
b.check_invariant()
k = min(max(i, 0), n)
assert a.to_list() == list(range(k))
assert b.to_list() == list(range(k, n))
c = a.append(b)
c.check_invariant()
assert c.to_list() == list(range(n))
for i in range(n):
assert t.lookup(i) == i
def test_random_against_lists() -> None:
rng = random.Random(20261104)
versions: list[tuple[FingerTree[int], list[int]]] = [(FingerTree(), [])]
checks = 0
for _ in range(3000):
t, ref = rng.choice(versions)
op = rng.randrange(7)
if op == 0:
x = rng.randrange(1000)
nt, nref = t.cons(x), [x] + ref
elif op == 1:
x = rng.randrange(1000)
nt, nref = t.snoc(x), ref + [x]
elif op == 2:
v = t.viewl()
if v is None:
assert ref == []
continue
assert v[0] == ref[0]
nt, nref = v[1], ref[1:]
elif op == 3:
w = t.viewr()
if w is None:
assert ref == []
continue
assert w[1] == ref[-1]
nt, nref = w[0], ref[:-1]
elif op == 4:
u, uref = rng.choice(versions)
nt, nref = t.append(u), ref + uref
elif op == 5:
i = rng.randrange(len(ref) + 1)
a, b = t.split_at(i)
a.check_invariant()
b.check_invariant()
assert a.to_list() == ref[:i] and b.to_list() == ref[i:]
nt, nref = (a, ref[:i]) if rng.random() < 0.5 else (b, ref[i:])
else:
if ref:
i = rng.randrange(len(ref))
assert t.lookup(i) == ref[i]
continue
nt.check_invariant()
assert len(nt) == len(nref)
assert nt.to_list() == nref
checks += 1
if len(nref) < 3000:
versions.append((nt, nref))
for t, ref in versions:
assert t.to_list() == ref
assert checks >= 500
What the next level needs
- the tests did not run in this build