Position heap in Python
Trie over the reversed prefixes of a text that indexes the position of each pattern occurrence.
- File
position_heap.py, 82 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/position-heap/. The authors are listed in prov/PDF2MC.jsonl.
class PositionHeap:
"""The position heap of a text, built by prepending symbols.
The suffixes are inserted from the shortest to the longest, each as the
shortest prefix of it that is not yet a node. Every node but the root holds
one suffix, stored by its length so that prepending does not renumber it.
Seen from the reversed text, the suffixes are its reversed prefixes, and
prepending to the text is appending to the reversed text.
"""
def __init__(self, text: str = "") -> None:
self.rev: list[str] = [] # the text, reversed
self.children: list[dict[str, int]] = [{}] # node 0 is the root
self.length: list[int] = [0] # suffix length held by each node
self.depth: list[int] = [0]
for c in reversed(text):
self.prepend(c)
@property
def text(self) -> str:
return "".join(reversed(self.rev))
def _sym(self, ell: int, k: int) -> str:
"""Symbol k of the suffix of length ell."""
return self.rev[ell - 1 - k]
def prepend(self, c: str) -> None:
"""Make the text c + text and insert its new longest suffix."""
self.rev.append(c)
ell = len(self.rev)
x, k = 0, 0 # @line L1
while self._sym(ell, k) in self.children[x]: # @line L2
x = self.children[x][self._sym(ell, k)]
k += 1
y = len(self.length) # @line L3
self.children.append({})
self.length.append(ell)
self.depth.append(k + 1)
self.children[x][self._sym(ell, k)] = y
def _occurs_at(self, p: str, ell: int) -> bool:
return ell >= len(p) and all(self._sym(ell, k) == p[k] for k in range(len(p)))
def search(self, p: str) -> list[int]:
"""The sorted start positions of the occurrences of p in the text."""
if p == "":
raise ValueError("the pattern must be nonempty")
n = len(self.rev)
found: list[int] = []
x, k = 0, 0 # @line L4
while k < len(p) and p[k] in self.children[x]:
x = self.children[x][p[k]]
k += 1
if k < len(p) and self._occurs_at(p, self.length[x]): # @line L5
found.append(self.length[x])
if k == len(p): # @line L6
stack = [x]
while stack:
y = stack.pop()
found.append(self.length[y])
stack.extend(self.children[y].values())
return sorted(n - ell for ell in found) # @line L7
def check_invariant(self) -> None:
n = len(self.rev)
assert len(self.length) == n + 1 and sorted(self.length) == list(range(n + 1))
stack = [(0, "")]
seen = 0
while stack:
x, s = stack.pop()
seen += 1
assert self.depth[x] == len(s)
ell = self.length[x]
assert ell >= len(s) and all(self._sym(ell, k) == s[k] for k in range(len(s)))
for c, y in self.children[x].items():
assert self.length[y] > ell, "heap order"
stack.append((y, s + c))
assert seen == n + 1
test_position_heap.py 65 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/position-heap/. The authors are listed in prov/PDF2MC.jsonl.
import random
import pytest
from position_heap import PositionHeap
def occurrences(t: str, p: str) -> list[int]:
return [i for i in range(len(t) - len(p) + 1) if t.startswith(p, i)]
def test_example() -> None:
h = PositionHeap("abaabab")
h.check_invariant()
n = 7
nodes = {}
stack = [(0, "")]
while stack:
x, s = stack.pop()
if x:
nodes[s] = n - h.length[x]
stack.extend((y, s + c) for c, y in h.children[x].items())
assert nodes == {"a": 5, "aa": 2, "ab": 3, "aba": 0, "b": 6, "ba": 4, "baa": 1}
assert h.search("ab") == [0, 3, 5]
assert h.search("abab") == [3]
assert h.search("b") == [1, 4, 6]
assert h.search("bb") == []
def test_edges() -> None:
h = PositionHeap()
h.check_invariant()
assert h.search("a") == []
h = PositionHeap("aaaaa")
h.check_invariant()
assert max(h.depth) == 5
assert h.search("aa") == [0, 1, 2, 3]
assert h.search("aaaaaa") == []
with pytest.raises(ValueError):
h.search("")
def test_prepend_online() -> None:
h = PositionHeap()
t = ""
for c in "mississippi"[::-1]:
h.prepend(c)
t = c + t
h.check_invariant()
assert h.text == t
assert h.search("ss") == occurrences(t, "ss")
def test_random_against_scan() -> None:
rng = random.Random(2011)
for _ in range(600):
t = "".join(rng.choice("ab" if rng.random() < 0.5 else "abc") for _ in range(rng.randint(0, 30)))
h = PositionHeap(t)
h.check_invariant()
for _ in range(4):
p = "".join(rng.choice("abc") for _ in range(rng.randint(1, 5)))
assert h.search(p) == occurrences(t, p)
What the next level needs
- the tests did not run in this build