Main-Lorentz algorithm in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/main-lorentz-algorithm/. The authors are listed in prov/W9BQDM.jsonl.

from typing import Hashable, List, Sequence, Tuple

_SEP = object()  # a separator equal to no letter


def z_function(t: Sequence[object]) -> List[int]:
    """z[i] = length of the longest common prefix of t and t[i:], with z[0] = len(t)."""
    n = len(t)
    z = [0] * n
    if n == 0:
        return z
    z[0] = n
    left = right = 0  # t[left:right] is a prefix of t, right maximal so far
    for i in range(1, n):
        k = min(right - i, z[i - left]) if i < right else 0
        while i + k < n and t[k] == t[i + k]:
            k += 1
        z[i] = k
        if i + k > right:
            left, right = i, i + k
    return z


def main_lorentz(s: Sequence[Hashable]) -> List[Tuple[int, int, int]]:
    """Return triples (lo, hi, m): s[i:i+2m] is a square for each lo <= i <= hi.

    Every occurrence of a square (a pair (i, m) with s[i:i+m] == s[i+m:i+2m],
    m >= 1) lies in exactly one triple. Letters are only tested for equality.
    Indices are 0-based. Cost: O(n log n) time, O(n) extra space besides the output.
    """
    n = len(s)
    # @line L1
    out: List[Tuple[int, int, int]] = []
    stack = [(0, n)]
    # @line L2
    while stack:
        # @line L3
        lo, hi = stack.pop()
        if hi - lo <= 1:
            continue
        # @line L4
        mid = (lo + hi) // 2
        u = list(s[lo:mid])
        v = list(s[mid:hi])
        nu, nv = len(u), len(v)
        stack.append((lo, mid))
        stack.append((mid, hi))
        # @line L5
        ru = u[::-1]
        rv = v[::-1]
        z_ru = z_function(ru)
        z_vu = z_function(v + [_SEP] + u)
        z_rurv = z_function(ru + [_SEP] + rv)
        z_v = z_function(v)
        # @line L6
        for m in range(1, nu + 1):
            # @line L7
            k1 = z_ru[m] if m < nu else 0          # common suffix of s[lo:mid-m] and s[lo:mid]
            k2 = z_vu[nv + 1 + nu - m]             # common prefix of s[mid-m:mid] and s[mid:hi]
            # @line L8
            a1 = max(0, m - k2)
            a2 = min(m - 1, k1)
            if a1 <= a2:
                out.append((mid - m - a2, mid - m - a1, m))
        # @line L9
        for m in range(1, nv + 1):
            # @line L10
            k1 = z_rurv[nu + 1 + nv - m]           # common suffix of s[lo:mid] and s[mid:mid+m]
            k2 = z_v[m] if m < nv else 0           # common prefix of s[mid:hi] and s[mid+m:hi]
            # @line L11
            a1 = max(1, m - k2)
            a2 = min(m - 1, k1)
            if a1 <= a2:
                out.append((mid - a2, mid - a1, m))
    # @line L12
    return out


def square_occurrences(s: Sequence[Hashable]) -> List[Tuple[int, int]]:
    """All pairs (i, m) with s[i:i+m] == s[i+m:i+2m], m >= 1, sorted."""
    pairs = [(i, m) for lo, hi, m in main_lorentz(s) for i in range(lo, hi + 1)]
    pairs.sort()
    return pairs


def is_square_free(s: Sequence[Hashable]) -> bool:
    """True if s has no factor uu with u nonempty."""
    return not main_lorentz(s)