Counting inversions by merge sort in Python

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

from typing import Protocol, Sequence, TypeVar


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


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


def sort_and_count(a: list[T], lo: int = 0, hi: int | None = None) -> int:
    """Return the number of inversions of a[lo:hi] and sort a[lo:hi] in place.

    An inversion is a pair p < q with a[p] > a[q]; equal items are not counted.
    Time O(n log n) and O(n) extra space for n = hi - lo. The recursion depth is
    about lg n, so the Python stack is not a problem.
    """
    if hi is None:
        hi = len(a)
    if hi - lo <= 1:  # @line L1
        return 0
    mid = lo + (hi - lo) // 2  # @line L2
    x = sort_and_count(a, lo, mid)  # @line L3
    y = sort_and_count(a, mid, hi)  # @line L4
    return x + y + merge_and_count(a, lo, mid, hi)  # @line L5


def merge_and_count(a: list[T], lo: int, mid: int, hi: int) -> int:
    """Merge the sorted parts a[lo:mid] and a[mid:hi]; return the pairs (left > right)."""
    b = a[lo:mid]  # @line L6
    i, j, k, c = 0, mid, lo, 0  # @line L7
    while i < len(b) and j < hi:  # @line L8
        if not a[j] < b[i]:  # @line L9: b[i] <= a[j], with only <
            a[k] = b[i]
            i += 1
        else:
            a[k] = a[j]  # @line L10
            j += 1
            c += len(b) - i
        k += 1  # @line L11
    while i < len(b):  # @line L12
        a[k] = b[i]  # @line L13
        i += 1
        k += 1
    return c  # @line L14


def count_inversions(items: Sequence[T]) -> int:
    """Return the number of inversions of items, without changing it."""
    return sort_and_count(list(items))


def smaller_to_the_right(items: Sequence[T]) -> list[int]:
    """For each index p, return the number of q > p with items[q] < items[p].

    This is the per-item variant of the page: sort (value, index) pairs and,
    when a left item is placed, add the number of right items already placed.
    """
    n = len(items)
    counts = [0] * n
    order = list(range(n))
    width = 1
    buf = [0] * n
    while width < n:  # bottom-up merges, so no recursion is needed
        for lo in range(0, n, 2 * width):
            mid = min(lo + width, n)
            hi = min(lo + 2 * width, n)
            i, j, k = lo, mid, lo
            while i < mid and j < hi:
                if not items[order[j]] < items[order[i]]:
                    counts[order[i]] += j - mid
                    buf[k] = order[i]
                    i += 1
                else:
                    buf[k] = order[j]
                    j += 1
                k += 1
            while i < mid:
                counts[order[i]] += j - mid
                buf[k] = order[i]
                i += 1
                k += 1
            while j < hi:
                buf[k] = order[j]
                j += 1
                k += 1
            order[lo:hi] = buf[lo:hi]
        width *= 2
    return counts