Cascade merge sort in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/cascade-merge-sort/. The authors are listed in prov/KG5YAF.jsonl.

import heapq
from collections import deque
from typing import Protocol, TypeVar


class Ordered(Protocol):
    def __lt__(self, other: "Ordered", /) -> bool: ...


T = TypeVar("T", bound=Ordered)


class _Head:
    """The head of the run from input d, at position j."""

    __slots__ = ("x", "d", "j")

    def __init__(self, x: T, d: int, j: int) -> None:
        self.x, self.d, self.j = x, d, j

    def __lt__(self, other: "_Head") -> bool:
        if self.x < other.x:
            return True
        if other.x < self.x:
            return False
        return self.d < other.d


def _merge(runs: list[list[T]]) -> list[T]:
    """Merge sorted runs with a heap. An empty run is a dummy run."""
    heap = [_Head(r[0], d, 0) for d, r in enumerate(runs) if r]
    heapq.heapify(heap)
    out: list[T] = []
    while heap:
        h = heap[0]
        out.append(h.x)
        if h.j + 1 < len(runs[h.d]):
            heapq.heapreplace(heap, _Head(runs[h.d][h.j + 1], h.d, h.j + 1))
        else:
            heapq.heappop(heap)
    return out


def cascade_distribution(s: int, p: int) -> tuple[int, list[int]]:
    """The smallest level l and the cascade distribution a_1 >= ... >= a_p with sum >= s."""
    level, a = 0, [1] + [0] * (p - 1)
    while sum(a) < s:
        a = [sum(a[:p - i]) for i in range(p)]
        level += 1
    return level, a


def cascade_merge_sort(f: list[T], m: int = 4, t: int = 4) -> int:
    """Sort f in place with a cascade merge on t simulated tapes and a memory of m items.

    A tape is a queue of runs; an empty run is a dummy run. Precondition:
    m >= 1 and t >= 3. The sort is not stable. Returns the number of passes,
    which is the level of the smallest cascade distribution with at least
    s = ceil(n / m) runs. Each pass moves almost all of the data, in merges
    of order t - 1, t - 2, ..., 2; the time in memory is O(n log n).
    """
    if m < 1 or t < 3:
        raise ValueError("need m >= 1 and t >= 3")
    p = t - 1
    runs = [sorted(f[i:i + m]) for i in range(0, len(f), m)]  # @line L1
    if len(runs) <= 1:
        f[:] = runs[0] if runs else []
        return 0
    level, a = cascade_distribution(len(runs), p)  # @line L2
    got = [0] * p
    tape_runs: list[list[list[T]]] = [[] for _ in range(p)]
    for r in runs:  # @line L3
        d = max(range(p), key=lambda i: (a[i] - got[i], -i))
        tape_runs[d].append(r)
        got[d] += 1
    tapes: list[deque[list[T]]] = []
    for d in range(p):  # @line L4
        dummies: list[list[T]] = [[] for _ in range(a[d] - got[d])]
        tapes.append(deque(dummies + tape_runs[d]))
    tapes.append(deque())
    out = p  # @line L5
    passes = 0
    while sum(len(q) for q in tapes) > 1:  # @line L6
        active = [d for d in range(t) if d != out]  # @line L7
        while len(active) >= 2:  # @line L8
            k = min(len(tapes[d]) for d in active)  # @line L9
            for _ in range(k):  # @line L10
                tapes[out].append(_merge([tapes[d].popleft() for d in active]))  # @line L11
            emptied = [d for d in active if not tapes[d]]  # @line L12
            active = [d for d in active if tapes[d]]
            out = emptied[0]
        passes += 1  # @line L13
    last = next(q for q in tapes if q)
    f[:] = last[0]  # @line L14
    assert passes == level
    return passes