Compressed suffix tree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/compressed-suffix-tree/. The authors are listed in prov/Q61TVH.jsonl.

from collections.abc import Hashable, Sequence

INF = 1 << 62
Node = tuple[int, int]


def _suffix_array(t: list[int]) -> list[int]:
    """Suffix array by prefix doubling with library sorts; t ends with a unique smallest code."""
    n = len(t)
    sa, rank, h = list(range(n)), t[:], 1
    while n > 1:
        key = [(rank[i], rank[i + h] if i + h < n else -1) for i in range(n)]
        sa.sort(key=key.__getitem__)
        new = [0] * n
        for r in range(1, n):
            new[sa[r]] = new[sa[r - 1]] + (key[sa[r]] != key[sa[r - 1]])
        rank = new
        if rank[sa[-1]] == n - 1:
            break
        h *= 2
    return sa


class _GammaStream:
    """Positive integers written with the Elias gamma code into one growing integer."""

    def __init__(self) -> None:
        self.bits, self.length = 0, 0

    def write(self, x: int) -> None:
        k = x.bit_length()
        self.length += k - 1
        self.bits |= int(format(x, "b")[::-1], 2) << self.length
        self.length += k

    def read(self, pos: int) -> tuple[int, int]:
        zeros = 0
        while not self.bits >> (pos + zeros) & 1:
            zeros += 1
        pos += zeros
        chunk = (self.bits >> pos) & ((1 << (zeros + 1)) - 1)
        return int(format(chunk, "0%db" % (zeros + 1))[::-1], 2), pos + zeros + 1


class CompressedSuffixTree:
    """Suffix tree of T$ from a compressed suffix array, the PLCP bit vector H and range minima.

    The symbols need a total order and are coded 1 .. sigma; the sentinel $ is the
    code 0, unique and smallest. Ranks and positions are 0-based, N = n + 1. A node
    is a rank range (lb, rb): a leaf has lb == rb, the root is (0, n). Neither the
    text, nor SA, nor LCP is kept after the build: SA comes from Psi and samples,
    LCP[r] = PLCP[SA[r]] comes from H, and the minima of LCP are kept per block of
    `block` ranks in a small tree. The boundary values LCP[0] and LCP[N] are -1.
    """

    def __init__(self, text: Sequence[Hashable], step: int = 4, block: int = 8) -> None:
        self.n = n = len(text)
        self.N, self.step, self.block = n + 1, step, block
        self.symbols = sorted(set(text))
        code = {c: r + 1 for r, c in enumerate(self.symbols)}
        t = [code[c] for c in text] + [0]
        sa = _suffix_array(t)
        isa = [0] * (n + 1)
        for i, p in enumerate(sa):
            isa[p] = i
        # Psi as the increasing values F(i) * N + Psi(i), sampled and gamma coded.
        self.samples: list[tuple[int, int]] = []
        self.stream, prev = _GammaStream(), 0
        for i, p in enumerate(sa):
            value = t[p] * self.N + isa[(p + 1) % self.N]
            if i % block == 0:
                self.samples.append((value, self.stream.length))
            else:
                self.stream.write(value - prev)
            prev = value
        self.sa_sample = {i: p for i, p in enumerate(sa) if p % step == 0 or p == n}
        self.isa_sample = [isa[p] for p in range(0, self.N, step)]
        # PLCP by Kasai's method, then H with a 1 at PLCP[p] + 2p.
        plcp, h = [0] * self.N, 0
        for p in range(self.N):
            if isa[p] > 0:
                q = sa[isa[p] - 1]
                while t[p + h] == t[q + h]:
                    h += 1
                plcp[p] = h
                h = max(h - 1, 0)
            else:
                h = 0
        self.h_bits, self.h_samples = 0, []
        for p in range(self.N):
            if p % block == 0:
                self.h_samples.append(plcp[p] + 2 * p)
            self.h_bits |= 1 << (plcp[p] + 2 * p)
        # Block minima of LCP[0 .. N] in a tree of minima.
        lcp = [-1] + [plcp[sa[r]] for r in range(1, self.N)] + [-1]
        nb = (self.N + 1 + block - 1) // block
        self.size = 1
        while self.size < nb:
            self.size *= 2
        self.tree = [INF] * (2 * self.size)
        for b in range(nb):
            self.tree[self.size + b] = min(lcp[b * block:(b + 1) * block])
        for v in range(self.size - 1, 0, -1):
            self.tree[v] = min(self.tree[2 * v], self.tree[2 * v + 1])

    # ----- the compressed suffix array -----

    def _value(self, i: int) -> int:
        value, pos = self.samples[i // self.block]
        for _ in range(i % self.block):
            gap, pos = self.stream.read(pos)
            value += gap
        return value

    def psi(self, i: int) -> int:
        return self._value(i) % self.N

    def first(self, i: int) -> int:
        return self._value(i) // self.N

    def sa(self, i: int) -> int:
        t = 0
        while i not in self.sa_sample:
            i, t = self.psi(i), t + 1
        return self.sa_sample[i] - t

    def isa(self, p: int) -> int:
        i = self.isa_sample[p // self.step]
        for _ in range(p % self.step):
            i = self.psi(i)
        return i

    # ----- LCP from H, and range minima -----

    def _plcp(self, p: int) -> int:
        """PLCP[p] = select1(H, p) - 2p, by a scan from the sampled one."""
        pos, k = self.h_samples[p // self.block], p % self.block
        while k:
            pos += 1
            if self.h_bits >> pos & 1:
                k -= 1
        return pos - 2 * p

    def lcp(self, r: int) -> int:
        """LCP[r] for 0 <= r <= N, with LCP[0] = LCP[N] = -1."""
        return -1 if r <= 0 or r >= self.N else self._plcp(self.sa(r))

    def _range_min(self, a: int, b: int) -> int:
        """Minimum of LCP[a .. b]."""
        ba, bb, B = a // self.block, b // self.block, self.block
        if ba == bb:
            return min(self.lcp(r) for r in range(a, b + 1))
        best = min(min(self.lcp(r) for r in range(a, (ba + 1) * B)), min(self.lcp(r) for r in range(bb * B, b + 1)))
        lo, hi = ba + 1 + self.size, bb - 1 + self.size
        while lo <= hi:
            if lo & 1:
                best, lo = min(best, self.tree[lo]), lo + 1
            if not hi & 1:
                best, hi = min(best, self.tree[hi]), hi - 1
            lo, hi = lo >> 1, hi >> 1
        return best

    def _fwd(self, start: int, x: int) -> int:
        """The smallest r >= start with LCP[r] < x; N is a valid answer for x >= 0."""
        B = self.block
        for r in range(start, min((start // B + 1) * B, self.N + 1)):
            if self.lcp(r) < x:
                return r
        i = start // B + 1 + self.size
        while self.tree[i] >= x:
            while i & 1:
                i >>= 1
            i += 1
        while i < self.size:
            i = 2 * i if self.tree[2 * i] < x else 2 * i + 1
        return next(r for r in range((i - self.size) * B, self.N + 1) if self.lcp(r) < x)

    def _bwd(self, start: int, x: int) -> int:
        """The largest r <= start with LCP[r] < x; 0 is a valid answer for x >= 0."""
        B = self.block
        for r in range(start, (start // B) * B - 1, -1):
            if self.lcp(r) < x:
                return r
        i = start // B - 1 + self.size
        while self.tree[i] >= x:
            while not i & 1:
                i >>= 1
            i -= 1
        while i < self.size:
            i = 2 * i + 1 if self.tree[2 * i + 1] < x else 2 * i
        b = i - self.size
        return next(r for r in range(min((b + 1) * B, self.N + 1) - 1, -1, -1) if self.lcp(r) < x)

    def rmq(self, a: int, b: int) -> int:
        """The leftmost position of the minimum of LCP[a .. b]."""
        return self._fwd(a, self._range_min(a, b) + 1)

    # ----- tree operations -----

    def root(self) -> Node:
        return (0, self.n)

    def is_leaf(self, v: Node) -> bool:
        return v[0] == v[1]

    def depth(self, v: Node) -> int:
        """String depth; a leaf counts the sentinel."""
        lb, rb = v
        if lb == rb:
            return self.N - self.sa(lb)  # @line L1
        return self.lcp(self.rmq(lb + 1, rb))  # @line L2

    def _enclose(self, k: int) -> Node:
        """The node whose value is LCP[k] and whose range contains k - 1 and k."""
        h = self.lcp(k)
        return (self._bwd(k - 1, h), self._fwd(k + 1, h) - 1)

    def parent(self, v: Node) -> Node | None:
        lb, rb = v
        if v == self.root():
            return None
        k = lb if self.lcp(lb) > self.lcp(rb + 1) else rb + 1  # @line L3
        return self._enclose(k)  # @line L4

    def first_child(self, v: Node) -> Node | None:
        lb, rb = v
        if lb == rb:
            return None
        return (lb, self.rmq(lb + 1, rb) - 1)  # @line L5

    def next_sibling(self, v: Node) -> Node | None:
        lb, rb = v
        if v == self.root() or self.lcp(lb) > self.lcp(rb + 1):  # @line L6
            return None
        h = self.lcp(rb + 1)  # @line L7
        return (rb + 1, self._fwd(rb + 2, h + 1) - 1)

    def children(self, v: Node) -> list[Node]:
        out, w = [], self.first_child(v)
        while w is not None:
            out.append(w)
            w = self.next_sibling(w)
        return out

    def lca(self, u: Node, v: Node) -> Node:
        if u[0] <= v[0] and v[1] <= u[1]:
            return u
        if v[0] <= u[0] and u[1] <= v[1]:
            return v
        if u[0] > v[0]:
            u, v = v, u
        return self._enclose(self.rmq(u[1] + 1, v[0]))  # @line L8

    def suffix_link(self, v: Node) -> Node | None:
        lb, rb = v
        if v == self.root():
            return None
        if lb == rb:
            return self.root() if self.sa(lb) == self.n else (self.psi(lb),) * 2
        return self.lca((self.psi(lb),) * 2, (self.psi(rb),) * 2)  # @line L9

    def letter(self, v: Node, d: int) -> int:
        """Code of the symbol at depth d of the path label of v (0 is the sentinel)."""
        return self.first(self.isa(self.sa(v[0]) + d))  # @line L10

    def child(self, v: Node, c: Hashable) -> Node | None:
        code = {s: r + 1 for r, s in enumerate(self.symbols)}.get(c)
        d = self.depth(v)
        for w in self.children(v):
            if self.letter(w, d) == code:  # @line L11
                return w
        return None

    def find(self, pattern: Sequence[Hashable]) -> tuple[int, int]:
        """Rank range [lo, hi) of the suffixes that start with pattern."""
        code = {s: r + 1 for r, s in enumerate(self.symbols)}
        if any(c not in code for c in pattern):
            return (0, 0)
        p = [code[c] for c in pattern]
        v, d = self.root(), 0
        while d < len(p):
            if self.is_leaf(v):
                return (0, 0)
            w = next((w for w in self.children(v) if self.letter(w, d) == p[d]), None)
            if w is None:
                return (0, 0)
            e = min(self.depth(w), len(p))
            i = self.isa(self.sa(w[0]) + d + 1) if d + 1 < e else 0
            for x in p[d + 1:e]:
                if self.first(i) != x:
                    return (0, 0)
                i = self.psi(i)
            v, d = w, e
        return (v[0], v[1] + 1)

    def label(self, v: Node) -> list[Hashable]:
        """The path label of v, without the sentinel."""
        out, d = [], self.depth(v)
        i = self.isa(self.sa(v[0])) if d else 0
        for _ in range(d):
            c = self.first(i)
            if c == 0:
                break
            out.append(self.symbols[c - 1])
            i = self.psi(i)
        return out