Permutation rank and unrank algorithm in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/permutation-rank-unrank/. The authors are listed in prov/NW945N.jsonl.

from typing import Sequence


def _full_tree(n: int) -> list[int]:
    """Fenwick tree (1-based) in which every value 1..n is present once."""
    return [0] + [i & -i for i in range(1, n + 1)]


def _prefix(tree: list[int], i: int) -> int:
    """Number of present values among 1..i."""
    s = 0
    while i > 0:  # @line L16
        s += tree[i]
        i -= i & -i
    return s  # @line L17


def _add(tree: list[int], n: int, i: int, d: int) -> None:
    """Add d to the count of value i."""
    while i <= n:  # @line L18
        tree[i] += d
        i += i & -i


def _select(tree: list[int], n: int, k: int) -> int:
    """Return the k-th smallest present value, for 1 <= k <= number present."""
    pos = 0  # @line L13
    step = 1 << (n.bit_length() - 1) if n > 0 else 0
    while step > 0:  # @line L14
        if pos + step <= n and tree[pos + step] < k:
            pos += step
            k -= tree[pos]
        step >>= 1
    return pos + 1  # @line L15


def rank(a: Sequence[int]) -> int:
    """Return the lexicographic rank (from 0) of a permutation a of 1..n.

    Time O(n log n) arithmetic operations; the rank is an exact Python int.
    """
    n = len(a)
    if sorted(a) != list(range(1, n + 1)):
        raise ValueError("a must be a permutation of 1..n")
    tree = _full_tree(n)  # @line L1
    r = 0
    for i in range(1, n + 1):  # @line L2
        c = _prefix(tree, a[i - 1] - 1)  # @line L3
        _add(tree, n, a[i - 1], -1)  # @line L4
        r = r * (n - i + 1) + c  # @line L5
    return r  # @line L6


def unrank(r: int, n: int) -> list[int]:
    """Return the permutation of 1..n with lexicographic rank r, for 0 <= r < n!."""
    if n < 0 or r < 0:
        raise ValueError("need n >= 0 and r >= 0")
    d = [0] * (n + 1)
    for i in range(n, 0, -1):  # @line L7
        d[i] = r % (n - i + 1)
        r //= n - i + 1
    if r != 0:
        raise ValueError("r must be less than n!")
    tree = _full_tree(n)  # @line L8
    a = [0] * n
    for i in range(1, n + 1):  # @line L9
        a[i - 1] = _select(tree, n, d[i] + 1)  # @line L10
        _add(tree, n, a[i - 1], -1)  # @line L11
    return a  # @line L12


def lehmer_code(a: Sequence[int]) -> list[int]:
    """Return c with c[i] = number of j > i with a[j] < a[i], in O(n log n) time."""
    n = len(a)
    tree = _full_tree(n)
    out = []
    for x in a:
        out.append(_prefix(tree, x - 1))
        _add(tree, n, x, -1)
    return out