Amir-Lewenstein-Porat algorithm in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/amir-lewenstein-porat-algorithm/. The authors are listed in prov/EGHY1F.jsonl.

"""The Amir-Lewenstein-Porat algorithm for the k-mismatch problem.

Edit model: Hamming distance (substitutions only). Indices are 0-based.
The result is the sorted list of start positions i in 0..len(text)-len(pattern)
such that pattern and text[i:i+len(pattern)] differ in at most k places.

Symbols that occur at least f = ceil(sqrt(k * ceil(log2 m))) times in the
pattern are frequent. With many frequent symbols, 2k marks of the pattern
filter the alignments and a kangaroo check with longest common extensions
verifies the few that remain. With few frequent symbols, one convolution per
frequent symbol and direct counting for the other symbols give every
Hamming distance. Convolutions use a number-theoretic transform modulo the
prime 998244353, exact because every count is at most m. All loops are
iterative.
"""

from math import isqrt

MOD = 998244353
ROOT = 3  # a primitive root modulo MOD


def ntt(a: list[int], invert: bool) -> None:
    """Transform a in place; len(a) is a power of two that divides 2**23."""
    n = len(a)
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j |= bit
        if i < j:
            a[i], a[j] = a[j], a[i]
    length = 2
    while length <= n:
        w = pow(ROOT, (MOD - 1) // length, MOD)
        if invert:
            w = pow(w, MOD - 2, MOD)
        half = length >> 1
        for start in range(0, n, length):
            wn = 1
            for k in range(start, start + half):
                u = a[k]
                v = a[k + half] * wn % MOD
                a[k] = (u + v) % MOD
                a[k + half] = (u - v) % MOD
                wn = wn * w % MOD
        length <<= 1
    if invert:
        inv_n = pow(n, MOD - 2, MOD)
        for i in range(n):
            a[i] = a[i] * inv_n % MOD


class LCE:
    """Longest common extension of a suffix of a and a suffix of b."""

    def __init__(self, a: str, b: str) -> None:
        syms = sorted(set(a) | set(b))
        code = {c: r + 2 for r, c in enumerate(syms)}
        s = [code[c] for c in a] + [1] + [code[c] for c in b] + [0]
        self.off = len(a) + 1
        n = len(s)
        rank = s[:]
        sa = list(range(n))
        h = 1
        while True:  # prefix doubling: sort by (rank[i], rank[i + h])
            key = [rank[i] * (n + 2) + (rank[i + h] + 1 if i + h < n else 0) for i in range(n)]
            sa.sort(key=key.__getitem__)
            new = [0] * n
            for t in range(1, n):
                new[sa[t]] = new[sa[t - 1]] + (key[sa[t]] != key[sa[t - 1]])
            rank = new
            if rank[sa[-1]] == n - 1:
                break
            h *= 2
        self.rank = rank
        lcp = [0] * n  # lcp[t] = lcp of suffixes sa[t - 1] and sa[t]
        q = 0
        for i in range(n):
            if rank[i] == 0:
                q = 0
                continue
            j = sa[rank[i] - 1]
            while i + q < n and j + q < n and s[i + q] == s[j + q]:
                q += 1
            lcp[rank[i]] = q
            if q:
                q -= 1
        self.table = [lcp]
        w = 1
        while 2 * w <= n:
            prev = self.table[-1]
            self.table.append([min(prev[t], prev[t + w]) for t in range(n - 2 * w + 1)])
            w *= 2

    def query(self, i: int, j: int) -> int:
        """Length of the longest common prefix of a[i:] and b[j:]."""
        x, y = self.rank[i], self.rank[self.off + j]
        if x > y:
            x, y = y, x
        x += 1  # minimum of lcp[x..y]
        lev = (y - x + 1).bit_length() - 1
        row = self.table[lev]
        return min(row[x], row[y - (1 << lev) + 1])


def parameters(m: int, k: int) -> tuple[int, int]:
    """Return (f, g): the frequency threshold and the number of mark symbols."""
    lam = (m - 1).bit_length() if m > 1 else 0  # ceil(log2 m)
    x = k * lam
    f = isqrt(x)
    if f * f < x:
        f += 1  # f = ceil(sqrt(k * lam))
    f = max(1, f)
    return f, -(-2 * k // f)  # g = ceil(2k / f)


def kangaroo(lce: LCE, m: int, i: int, k: int) -> bool:
    """Return True when the window at i has at most k mismatches."""
    q, e = 0, 0
    while q < m:
        q += lce.query(q, i + q)  # jump over a run of equal symbols
        if q < m:
            e += 1
            if e > k:
                return False
            q += 1
    return True


def correlate(pattern: str, text: str, c: str, out: list[int]) -> None:
    """Add to out[i] the number of q with pattern[q] = text[i + q] = c."""
    m, n = len(pattern), len(text)
    size = 1
    while size < 3 * m - 2:
        size <<= 1
    x = [0] * size
    for q in range(m):
        if pattern[m - 1 - q] == c:
            x[q] = 1
    ntt(x, False)
    for b in range(0, n - m + 1, m):  # text blocks of length 2m - 1
        y = [0] * size
        for j, t in enumerate(text[b:b + 2 * m - 1]):
            if t == c:
                y[j] = 1
        ntt(y, False)
        z = [u * v % MOD for u, v in zip(x, y)]
        ntt(z, True)
        for t in range(min(m, n - m + 1 - b)):
            out[b + t] += z[t + m - 1]


def alp_k_mismatch(pattern: str, text: str, k: int) -> list[int]:
    """Return the start positions with at most k mismatches.

    Precondition: k >= 0. Cost: O(n + n * sqrt(k log m)) time after the
    longest common extension structure, O(n + m) extra space.
    """
    if k < 0:
        raise ValueError("k must be at least 0")
    m, n = len(pattern), len(text)
    if m > n:
        return []
    if k >= m:  # @line L1
        return list(range(n - m + 1))
    f, g = parameters(m, k)  # @line L2
    occ: dict[str, list[int]] = {}
    for q, c in enumerate(pattern):
        occ.setdefault(c, []).append(q)
    freq = sorted(c for c in occ if len(occ[c]) >= f)
    if len(freq) >= g:  # @line L3
        marks = {c: occ[c][:f] for c in freq[:g]}  # @line L4
        count = [0] * (n - m + 1)
        for j, t in enumerate(text):  # @line L5
            for q in marks.get(t, ()):
                if 0 <= j - q <= n - m:
                    count[j - q] += 1
        lce = LCE(pattern, text)
        need = g * f - k
        return [i for i in range(n - m + 1) if count[i] >= need and kangaroo(lce, m, i, k)]  # @line L6
    match = [0] * (n - m + 1)  # @line L7
    for c in freq:  # @line L8
        correlate(pattern, text, c, match)
    frequent = set(freq)
    for j, t in enumerate(text):  # @line L9
        if t not in frequent:
            for q in occ.get(t, ()):
                if 0 <= j - q <= n - m:
                    match[j - q] += 1
    return [i for i in range(n - m + 1) if m - match[i] <= k]  # @line L10


def case_of(pattern: str, k: int) -> str:
    """Return "A" (marks and kangaroo) or "B" (convolutions), for tests."""
    m = len(pattern)
    f, g = parameters(m, k)
    counts: dict[str, int] = {}
    for c in pattern:
        counts[c] = counts.get(c, 0) + 1
    return "A" if sum(1 for v in counts.values() if v >= f) >= g else "B"