Mo's algorithm with updates in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/mo-algorithm-with-updates/. The authors are listed in prov/XFPAYN.jsonl.

from typing import Callable, Hashable, Sequence, TypeVar

T = TypeVar("T")
V = TypeVar("V")


def default_block(n: int) -> int:
    """The block size max(1, ceil(n^(2/3)))."""
    b = 1
    while b * b * b < n * n:
        b += 1
    return max(1, b)


def mo_with_updates(
    a: Sequence[V],
    queries: Sequence[tuple[int, int, int]],
    updates: Sequence[tuple[int, V]],
    add: Callable[[V], None],
    remove: Callable[[V], None],
    answer: Callable[[], T],
    block: int = 0,
) -> list[T]:
    """Answer offline window queries (l, r, t) on the array after t updates.

    A query (l, r, t) asks about a[l:r] after the first t updates of
    updates have been applied; an update (p, x) sets a[p] = x. add(x) and
    remove(x) put the value x into the window or take it out, and answer()
    reads the window. block <= 0 picks ceil(n^(2/3)).
    Cost: O(q B + n^2 / B + m n^2 / B^2) calls of add and remove for q
    queries and m updates, plus O(q log q) to sort.
    """
    arr = list(a)
    upd = [list(u) for u in updates]
    n, q = len(arr), len(queries)
    b = block if block > 0 else default_block(n)
    order = sorted(range(q), key=lambda j: (queries[j][0] // b, queries[j][1] // b, queries[j][2]))  # @line L2
    cur_l = cur_r = cur_t = 0  # @line L3
    ans: list[T] = [None] * q  # type: ignore[list-item]

    def apply(j: int) -> None:
        p = upd[j][0]  # @line L10
        inside = cur_l <= p < cur_r
        if inside:  # @line L11
            remove(arr[p])
        arr[p], upd[j][1] = upd[j][1], arr[p]  # @line L12
        if inside:  # @line L13
            add(arr[p])

    for j in order:  # @line L4
        l, r, t = queries[j]
        while cur_t < t:  # @line L5
            apply(cur_t)
            cur_t += 1
        while cur_t > t:  # @line L6
            cur_t -= 1
            apply(cur_t)
        while cur_l > l:  # @line L7
            cur_l -= 1
            add(arr[cur_l])
        while cur_r < r:
            add(arr[cur_r])
            cur_r += 1
        while cur_l < l:
            remove(arr[cur_l])
            cur_l += 1
        while cur_r > r:
            cur_r -= 1
            remove(arr[cur_r])
        ans[j] = answer()  # @line L8
    return ans  # @line L9


def distinct_counts_with_updates(a: Sequence[Hashable], ops: Sequence[tuple[str, int, object]]) -> list[int]:
    """Answers of ("Q", l, r) operations, the number of distinct values of
    a[l:r], interleaved with ("U", p, x) operations that set a[p] = x."""
    queries: list[tuple[int, int, int]] = []
    updates: list[tuple[int, Hashable]] = []
    for kind, u, w in ops:  # @line L1
        if kind == "U":
            if not 0 <= u < len(a):
                raise IndexError(f"update position {u} out of range")
            updates.append((u, w))  # type: ignore[arg-type]
        else:
            r = int(w)  # type: ignore[call-overload]
            if not 0 <= u <= r <= len(a):
                raise IndexError(f"bad range [{u}, {r})")
            queries.append((u, r, len(updates)))
    count: dict[Hashable, int] = {}
    distinct = [0]

    def add(x: Hashable) -> None:
        c = count.get(x, 0)
        if c == 0:
            distinct[0] += 1
        count[x] = c + 1

    def remove(x: Hashable) -> None:
        c = count[x] - 1
        count[x] = c
        if c == 0:
            distinct[0] -= 1

    return mo_with_updates(a, queries, updates, add, remove, lambda: distinct[0])