Introselect in Python

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

from typing import Protocol, Sequence, TypeVar


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


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


def median3(x: T, y: T, z: T) -> T:
    """Return the median of three items, with at most three comparisons."""
    if y < x:
        x, y = y, x
    if z < y:
        y = z
        if y < x:
            y = x
    return y


def partition3(b: list[T], lo: int, hi: int, x: T) -> tuple[int, int]:
    """Rearrange b[lo .. hi - 1] into items < x, items equal to x, items > x.

    Returns (p, q) with b[lo .. p - 1] < x, b[p .. q - 1] equal to x, and
    b[q .. hi - 1] > x. This is the Dutch national flag partition.
    """
    p, i, q = lo, lo, hi
    while i < q:
        if b[i] < x:
            b[p], b[i] = b[i], b[p]
            p += 1
            i += 1
        elif x < b[i]:
            q -= 1
            b[i], b[q] = b[q], b[i]
        else:
            i += 1
    return p, q


def mom_select(s: list[T], k: int) -> T:
    """Return the item of rank k (1-based) of s by median of medians, in O(n) time."""
    while True:
        n = len(s)
        if n <= 5:
            return sorted(s)[k - 1]
        medians = [sorted(s[i:i + 5])[(min(5, n - i) - 1) // 2] for i in range(0, n, 5)]
        x = mom_select(medians, (len(medians) + 1) // 2)
        less = [y for y in s if y < x]
        greater = [y for y in s if x < y]
        equal = n - len(less) - len(greater)
        if k <= len(less):
            s = less
        elif k <= len(less) + equal:
            return x
        else:
            s, k = greater, k - len(less) - equal


def introselect(a: Sequence[T], k: int, limit: int = 2) -> T:
    """Return the item of rank k (1-based) of a, in O(n) time in the worst case.

    Runs quickselect with a median-of-three pivot and three-way partition.
    After limit partitions in a row that keep more than 3/4 of the range, it finishes
    the current range with median of medians. The input is not changed.
    """
    n = len(a)
    if not 1 <= k <= n:
        raise ValueError("k must be between 1 and len(a)")
    b = list(a)
    lo, hi, t, bad = 0, n, k - 1, 0  # @line L1
    while hi - lo > 1:  # @line L2
        if bad >= limit:  # @line L3
            return mom_select(b[lo:hi], t - lo + 1)  # @line L3
        x = median3(b[lo], b[(lo + hi - 1) // 2], b[hi - 1])  # @line L4
        p, q = partition3(b, lo, hi, x)  # @line L5
        if p <= t < q:  # @line L6
            return x  # @line L6
        old = hi - lo  # @line L7
        if t < p:  # @line L8
            hi = p  # @line L8
        else:
            lo = q  # @line L8
        if 4 * (hi - lo) > 3 * old:  # @line L9
            bad += 1  # @line L9
        else:
            bad = 0  # @line L9
    return b[lo]  # @line L10