Ternary search tree in Python
Trie node with three children (less, equal, greater) that stores one character per node.
- File
ternary_search_tree.py, 247 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/ternary-search-tree/. The authors are listed in prov/BVWSE3.jsonl.
from typing import Generic, TypeVar
V = TypeVar("V")
class _Node(Generic[V]):
__slots__ = ("c", "left", "mid", "right", "val")
def __init__(self, c: str) -> None:
self.c = c
self.left: _Node[V] | None = None
self.mid: _Node[V] | None = None
self.right: _Node[V] | None = None
self.val: V | None = None
class TernarySearchTree(Generic[V]):
"""A ternary search tree (TST) with string keys of length at least 1.
Each node holds one symbol and three links. A search compares each symbol
of the key with one node on the middle path and with the nodes of the
left and right links that it passes, so get and put take O(m + d) symbol
comparisons, where d is the number of left and right steps. Values must
not be None. The empty string is not a key.
"""
def __init__(self) -> None:
self.root: _Node[V] | None = None
self.n = 0
def size(self) -> int:
return self.n
def _find(self, s: str) -> "_Node[V] | None":
x = self.root
d = 0
while x is not None:
if s[d] < x.c:
x = x.left
elif s[d] > x.c:
x = x.right
elif d < len(s) - 1:
x = x.mid
d += 1
else:
return x
return None
def get(self, s: str) -> V | None:
"""Return the value of key s, or None."""
if not s:
raise ValueError("keys must not be empty")
x = self.root # @line L1
d = 0
while x is not None: # @line L2
if s[d] < x.c: # @line L3
x = x.left
elif s[d] > x.c: # @line L4
x = x.right
elif d < len(s) - 1: # @line L5
x = x.mid
d += 1
else:
return x.val # @line L6
return None # @line L7
def put(self, s: str, v: V) -> None:
"""Set the value of key s to v (s nonempty, v not None)."""
if not s:
raise ValueError("keys must not be empty")
if v is None:
raise ValueError("a value must not be None")
if self.root is None:
self.root = _Node(s[0])
x = self.root # @line L8
d = 0
while True: # @line L9
c = s[d]
if c < x.c: # @line L10
if x.left is None:
x.left = _Node(c)
x = x.left
elif c > x.c: # @line L11
if x.right is None:
x.right = _Node(c)
x = x.right
elif d < len(s) - 1: # @line L12
d += 1
if x.mid is None:
x.mid = _Node(s[d])
x = x.mid
else:
if x.val is None: # @line L13
self.n += 1
x.val = v # @line L14
return
def delete(self, s: str) -> None:
"""Remove key s if present, and prune the nodes that are no longer needed."""
if not s:
return
# path of (node, parent, which link of the parent leads to node)
path: list[tuple[_Node[V], _Node[V] | None, str]] = []
x, parent, link, d = self.root, None, "root", 0 # @line L15
while x is not None:
path.append((x, parent, link))
if s[d] < x.c:
x, parent, link = x.left, x, "left"
elif s[d] > x.c:
x, parent, link = x.right, x, "right"
elif d < len(s) - 1:
x, parent, link, d = x.mid, x, "mid", d + 1
else:
break
if x is None or x.val is None: # @line L16
return
x.val = None # @line L17
self.n -= 1
# walk back up: remove each node with no value and no middle child
for node, par, lk in reversed(path): # @line L18
if node.val is not None or node.mid is not None: # @line L19
return
repl = self._bst_remove_root(node) # @line L20
self._set_link(par, lk, repl)
@staticmethod
def _bst_remove_root(node: "_Node[V]") -> "_Node[V] | None":
"""Return the sibling BST rooted at node with node removed."""
if node.left is None:
return node.right
if node.right is None:
return node.left
# take the smallest node of the right subtree as the new root
p, m = node, node.right
while m.left is not None:
p, m = m, m.left
if p is not node:
p.left = m.right
m.right = node.right
m.left = node.left
return m
def _set_link(self, par: "_Node[V] | None", link: str, child: "_Node[V] | None") -> None:
if par is None:
self.root = child
else:
setattr(par, link, child)
def keys_with_prefix(self, p: str) -> list[str]:
"""Return the keys that start with p, in lexicographic order."""
out: list[str] = [] # @line L21
if p == "":
start, prefix = self.root, ""
else:
x = self._find(p)
if x is None:
return out
if x.val is not None:
out.append(p)
start, prefix = x.mid, p
# explicit stack: ("node", x, prefix) visits a subtree; ("emit", ...) reports a key
stack: list[tuple[_Node[V] | None, str, int]] = [(start, prefix, 0)]
while stack: # @line L22
y, t, state = stack.pop()
if y is None:
continue
if state == 0: # in-order: left, node, mid, right
stack.append((y.right, t, 0))
stack.append((y, t, 1))
stack.append((y.left, t, 0))
else:
if y.val is not None: # @line L23
out.append(t + y.c)
stack.append((y.mid, t + y.c, 0)) # @line L24
return out # @line L25
def keys(self) -> list[str]:
return self.keys_with_prefix("")
def longest_prefix_of(self, s: str) -> str | None:
"""Return the longest key that is a prefix of s, or None."""
x, d, best = self.root, 0, -1
while x is not None and d < len(s):
if s[d] < x.c:
x = x.left
elif s[d] > x.c:
x = x.right
else:
d += 1
if x.val is not None:
best = d
x = x.mid
return None if best < 0 else s[:best]
def keys_that_match(self, q: str) -> list[str]:
"""Return the keys that match q, where '.' matches any one symbol, in order."""
out: list[str] = []
if not q:
return out
stack: list[tuple[_Node[V] | None, int, str, int]] = [(self.root, 0, "", 0)]
while stack:
y, d, t, state = stack.pop()
if y is None:
continue
c = q[d]
if state == 0:
if c == "." or c > y.c:
stack.append((y.right, d, t, 0))
stack.append((y, d, t, 1))
if c == "." or c < y.c:
stack.append((y.left, d, t, 0))
elif c == "." or c == y.c:
if d == len(q) - 1:
if y.val is not None:
out.append(t + y.c)
else:
stack.append((y.mid, d + 1, t + y.c, 0))
return out
def node_count(self) -> int:
count, stack = 0, [self.root]
while stack:
y = stack.pop()
if y is not None:
count += 1
stack.extend((y.left, y.mid, y.right))
return count
def check_invariant(self) -> None:
"""Check the BST order of each sibling tree, no dead nodes, and the count."""
valued = 0
stack: list[tuple[_Node[V] | None, str | None, str | None]] = [(self.root, None, None)]
while stack:
y, lo, hi = stack.pop()
if y is None:
continue
assert (lo is None or lo < y.c) and (hi is None or y.c < hi)
assert y.val is not None or y.mid is not None, "dead node"
if y.val is not None:
valued += 1
stack.append((y.left, lo, y.c))
stack.append((y.right, y.c, hi))
stack.append((y.mid, None, None))
assert valued == self.n
test_ternary_search_tree.py 78 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/ternary-search-tree/. The authors are listed in prov/BVWSE3.jsonl.
import random
import re
import pytest
from ternary_search_tree import TernarySearchTree
WORDS = "she sells sea shells by the sea shore".split()
def test_example() -> None:
t: TernarySearchTree[int] = TernarySearchTree()
for i, w in enumerate(WORDS):
t.put(w, i)
t.check_invariant()
assert t.size() == 7 and t.node_count() == 19
assert t.root is not None and t.root.c == "s"
assert t.get("sea") == 6 and t.get("shy") is None and t.get("sh") is None
assert t.keys() == ["by", "sea", "sells", "she", "shells", "shore", "the"]
assert t.keys_with_prefix("sh") == ["she", "shells", "shore"]
assert t.keys_with_prefix("she") == ["she", "shells"]
assert t.keys_that_match(".he") == ["she", "the"]
assert t.longest_prefix_of("shellsort") == "shells"
t.delete("shells")
t.check_invariant()
assert t.node_count() == 16
t.delete("she")
t.check_invariant()
assert t.keys() == ["by", "sea", "sells", "shore", "the"]
def test_edges() -> None:
t: TernarySearchTree[int] = TernarySearchTree()
assert t.get("a") is None and t.keys() == [] and t.longest_prefix_of("a") is None
with pytest.raises(ValueError):
t.put("", 1)
t.put("a", 1)
t.delete("a")
t.check_invariant()
assert t.root is None and t.size() == 0
for w in ["m", "c", "x", "a", "e", "d"]:
t.put(w, 0)
t.delete("c") # a node with two BST children
t.check_invariant()
assert t.keys() == ["a", "d", "e", "m", "x"]
def test_random_against_dict() -> None:
rng = random.Random(9103)
for _ in range(500):
t: TernarySearchTree[int] = TernarySearchTree()
model: dict[str, int] = {}
for _ in range(rng.randint(0, 40)):
k = "".join(rng.choice("abcd") for _ in range(rng.randint(1, 4)))
if rng.random() < 0.6:
v = rng.randint(0, 99)
t.put(k, v)
model[k] = v
else:
t.delete(k)
model.pop(k, None)
t.check_invariant()
assert t.get(k) == model.get(k)
assert t.size() == len(model)
assert t.keys() == sorted(model)
prefixes = {k[:i] for k in model for i in range(1, len(k) + 1)}
assert t.node_count() == len(prefixes)
p = "".join(rng.choice("abcd") for _ in range(rng.randint(0, 2)))
assert t.keys_with_prefix(p) == sorted(k for k in model if k.startswith(p))
s = "".join(rng.choice("abcd") for _ in range(rng.randint(0, 5)))
cands = [k for k in model if s.startswith(k)]
assert t.longest_prefix_of(s) == (max(cands, key=len) if cands else None)
q = "".join(rng.choice("ab.") for _ in range(rng.randint(1, 3)))
rx = re.compile(q)
assert t.keys_that_match(q) == sorted(k for k in model if rx.fullmatch(k))
What the next level needs
- the tests did not run in this build