Radix select in Python

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

from typing import Optional, Sequence


def radix_select(a: Sequence[int], k: int, r: int = 8, w: Optional[int] = None) -> int:
    """Return the item of rank k (1-based) of a list of nonnegative integers.

    Looks at the keys r bits at a time, from the most significant digit of a
    w-bit key down, and keeps only the bucket that holds rank k. If w is not
    given, it is the bit length of the largest key. Uses no comparisons
    between keys.
    """
    n = len(a)
    if not 1 <= k <= n:
        raise ValueError("k must be between 1 and len(a)")
    if r < 1:
        raise ValueError("r must be at least 1")
    if any(x < 0 for x in a):
        raise ValueError("keys must be nonnegative")
    if w is None:
        w = max(max(a).bit_length(), 1)
    mask = (1 << r) - 1
    s = list(a)  # @line L1
    d = -(-w // r)  # @line L1
    for j in range(d - 1, -1, -1):  # @line L2
        if len(s) == 1:  # @line L3
            return s[0]  # @line L3
        shift = r * j
        count = [0] * (mask + 1)  # @line L4
        for x in s:  # @line L5
            count[(x >> shift) & mask] += 1  # @line L5
        b = 0  # @line L6
        while k > count[b]:  # @line L7
            k -= count[b]  # @line L7
            b += 1  # @line L7
        s = [x for x in s if (x >> shift) & mask == b]  # @line L8
    return s[0]  # @line L9