Algorithm Search⌘K

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
position_heap.pyRawGitHub
# 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

Read the full page about Position heap