Enhanced suffix array in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/enhanced-suffix-array/. The authors are listed in prov/VFC544.jsonl.

from collections.abc import Hashable, Sequence

NIL = -1


def suffix_array(codes: list[int]) -> list[int]:
    """Suffix array of an integer string by prefix doubling with library sorts."""
    n = len(codes)
    rank = list(codes)
    sa = list(range(n))
    h = 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


def lcp_array(codes: list[int], sa: list[int]) -> list[int]:
    """LCP[0] = 0 and LCP[i] = lcp of the suffixes at ranks i - 1 and i, by Kasai's method."""
    n = len(codes)
    rank = [0] * n
    for i, p in enumerate(sa):
        rank[p] = i
    lcp, h = [0] * n, 0
    for p in range(n):
        if rank[p] > 0:
            q = sa[rank[p] - 1]
            while p + h < n and q + h < n and codes[p + h] == codes[q + h]:
                h += 1
            lcp[rank[p]] = h
            if h:
                h -= 1
        else:
            h = 0
    return lcp


class EnhancedSuffixArray:
    """Suffix array, LCP array and child table of T$ (Abouelhoda, Kurtz and Ohlebusch).

    A sentinel $ smaller than every symbol is appended, so there are N = n + 1
    suffixes, ranks 0..n, and rank 0 is the suffix "$" at position n. Positions and
    ranks are 0-based. The alphabet is general ordered. L has N + 1 entries with
    L[0] = 0 and L[N] = -1. The fields up, down and next hold NIL when undefined.
    """

    def __init__(self, text: Sequence[Hashable]) -> None:
        self.text = text
        n = len(text)
        code = {c: r + 1 for r, c in enumerate(sorted(set(text)))}
        codes = [code[c] for c in text] + [0]
        self.sa = suffix_array(codes)
        self.lcp = lcp_array(codes, self.sa) + [-1]
        self.n, self.size = n, n + 1
        self.up = [NIL] * (n + 2)
        self.down = [NIL] * (n + 2)
        self.next = [NIL] * (n + 2)
        self._child_table()

    def _child_table(self) -> None:
        L, N = self.lcp, self.size
        last, stack = NIL, [0]  # @line L1
        for i in range(1, N + 1):  # @line L2
            while stack and L[i] < L[stack[-1]]:  # @line L3
                last = stack.pop()  # @line L4
                if stack and L[i] <= L[stack[-1]] and L[stack[-1]] != L[last]:  # @line L5
                    self.down[stack[-1]] = last
            if last != NIL:  # @line L6
                self.up[i], last = last, NIL
            stack.append(i)  # @line L7
        stack = [0]  # @line L8
        for i in range(1, N + 1):  # @line L9
            while stack and L[i] < L[stack[-1]]:  # @line L10
                stack.pop()
            if stack and L[i] == L[stack[-1]]:  # @line L11
                self.next[stack.pop()] = i
            stack.append(i)  # @line L12

    def first_index(self, i: int, j: int) -> int:
        """The first l-index of the lcp-interval [i..j], i < j."""
        if (i, j) == (0, self.n):
            return self.next[0]
        u = self.up[j + 1]
        return u if i < u <= j else self.down[i]

    def lcp_value(self, i: int, j: int) -> int:
        """The value l of the lcp-interval [i..j], i < j: the string depth of the node."""
        return 0 if (i, j) == (0, self.n) else self.lcp[self.first_index(i, j)]

    def children(self, i: int, j: int) -> list[tuple[int, int]]:
        """The child intervals of the lcp-interval [i..j], i < j, in rank order; leaves are [r..r]."""
        out = []
        k = self.first_index(i, j)  # @line L13 @line L14
        out.append((i, k - 1))  # @line L15
        while self.next[k] != NIL:  # @line L16
            out.append((k, self.next[k] - 1))
            k = self.next[k]
        out.append((k, j))  # @line L17
        return out

    def symbol(self, r: int, d: int) -> Hashable | None:
        """Symbol at depth d of the suffix at rank r; None stands for the sentinel."""
        p = self.sa[r] + d
        return self.text[p] if p < self.n else None

    def child(self, i: int, j: int, c: Hashable) -> tuple[int, int] | None:
        """The child interval of [i..j] whose suffixes have c at the depth of [i..j]."""
        d = self.lcp_value(i, j)
        for a, b in self.children(i, j):
            if self.symbol(a, d) == c:
                return (a, b)
        return None

    def find(self, pattern: Sequence[Hashable]) -> tuple[int, int]:
        """Rank range [lo, hi) of the suffixes of T$ that start with pattern, top down."""
        m, n = len(pattern), self.n
        i, j, d = 0, n, 0
        while d < m:
            if i == j:
                return (0, 0)
            hit = self.child(i, j, pattern[d])
            if hit is None:
                return (0, 0)
            i, j = hit
            ell = n - self.sa[i] if i == j else self.lcp_value(i, j)
            e, p = min(ell, m), self.sa[i]
            if any(self.text[p + t] != pattern[t] for t in range(d + 1, e)):
                return (0, 0)
            d = e
        return (i, j + 1)

    def intervals(self) -> list[tuple[int, int, int]]:
        """All lcp-intervals as (value, lb, rb), by a top-down walk with an explicit stack."""
        if self.n == 0:
            return [(0, 0, 0)]
        out, todo = [], [(0, self.n)]
        while todo:
            i, j = todo.pop()
            out.append((self.lcp_value(i, j), i, j))
            todo.extend((a, b) for a, b in self.children(i, j) if a < b)
        return sorted(out, key=lambda x: (x[1], -x[2], x[0]))