Mo's algorithm in Python

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

import math
from typing import Callable, Hashable, Sequence, TypeVar

T = TypeVar("T")


def default_block(n: int, q: int) -> int:
    """The block size max(1, min(n, ceil(n / sqrt(q))))."""
    if n <= 0 or q <= 0:
        return 1
    s = math.isqrt(q)
    if s * s < q:
        s += 1
    return max(1, min(n, -(-n // s)))


def mo_algorithm(
    n: int,
    queries: Sequence[tuple[int, int]],
    add: Callable[[int], None],
    remove: Callable[[int], None],
    answer: Callable[[], T],
    block: int = 0,
) -> list[T]:
    """Answer offline window queries [l, r) over positions 0..n-1.

    Precondition: 0 <= l <= r <= n for each query; add(i) puts position i
    into the window, remove(i) takes it out, and answer() reads the value
    of the current window. block <= 0 picks ceil(n / sqrt(q)).
    Cost: O(n sqrt(q)) calls of add and remove, plus O(q log q) to sort.
    """
    q = len(queries)
    b = block if block > 0 else default_block(n, q)
    order = sorted(range(q), key=lambda j: (queries[j][0] // b, queries[j][1]))  # @line L1
    cur_l = cur_r = 0  # @line L2
    ans: list[T] = [None] * q  # type: ignore[list-item]
    for j in order:  # @line L3
        l, r = queries[j]
        while cur_l > l:  # @line L4
            cur_l -= 1
            add(cur_l)
        while cur_r < r:  # @line L5
            add(cur_r)
            cur_r += 1
        while cur_l < l:  # @line L6
            remove(cur_l)
            cur_l += 1
        while cur_r > r:  # @line L7
            cur_r -= 1
            remove(cur_r)
        ans[j] = answer()  # @line L8
    return ans  # @line L9


def distinct_counts(a: Sequence[Hashable], queries: Sequence[tuple[int, int]], block: int = 0) -> list[int]:
    """The number of distinct values of a[l:r] for each query (l, r).

    Precondition: 0 <= l <= r <= len(a). Cost: O(n sqrt(q) + q log q)
    expected time with a hash map of counts, O(n + q) extra space.
    """
    cnt: dict[Hashable, int] = {}
    d = 0

    def add(i: int) -> None:
        nonlocal d
        c = cnt.get(a[i], 0)
        if c == 0:
            d += 1
        cnt[a[i]] = c + 1

    def remove(i: int) -> None:
        nonlocal d
        c = cnt[a[i]] - 1
        if c < 0:
            raise ValueError("remove of an item that is not in the window")
        if c == 0:
            d -= 1
        cnt[a[i]] = c

    return mo_algorithm(len(a), queries, add, remove, lambda: d, block)