Merge sort tree in Python

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

from bisect import bisect_left, bisect_right
from typing import Callable, Optional, Sequence


def _merge(x: list[int], y: list[int]) -> list[int]:
    """Merge two sorted lists into one sorted list in O(len(x) + len(y))."""
    out: list[int] = []
    i = j = 0
    while i < len(x) and j < len(y):
        if y[j] < x[i]:
            out.append(y[j])
            j += 1
        else:
            out.append(x[i])
            i += 1
    out.extend(x[i:])
    out.extend(y[j:])
    return out


class MergeSortTree:
    """Static array with counting, successor and k-th queries on ranges.

    Build: O(n log n) time and space. count_less, count_leq and successor:
    O(log^2 n). kth_smallest: O(log^3 n). Ranges are half-open [l, r).
    """

    def __init__(self, a: Sequence[int]) -> None:
        """Build the tree of sorted lists over a."""
        self.n = len(a)
        size = 1
        while size < self.n:
            size *= 2
        self.size = size
        self.lists: list[list[int]] = [[] for _ in range(2 * size)]
        for i, x in enumerate(a):  # @line L7
            self.lists[size + i] = [x]
        for v in range(size - 1, 0, -1):  # @line L8
            self.lists[v] = _merge(self.lists[2 * v], self.lists[2 * v + 1])

    def _check_range(self, l: int, r: int) -> None:
        if not 0 <= l <= r <= self.n:
            raise IndexError(f"bad range [{l}, {r}) for length {self.n}")

    def _count(self, l: int, r: int, search: Callable[[list[int]], int]) -> int:
        """Sum search(S_v) over the canonical nodes of [l, r)."""
        self._check_range(l, r)
        c = 0
        l += self.size  # @line L1
        r += self.size
        while l < r:  # @line L2
            if l & 1:  # @line L3
                c += search(self.lists[l])
                l += 1
            if r & 1:  # @line L4
                r -= 1
                c += search(self.lists[r])
            l //= 2  # @line L5
            r //= 2
        return c  # @line L6

    def count_less(self, l: int, r: int, x: int) -> int:
        """Number of k in [l, r) with a[k] < x."""
        return self._count(l, r, lambda s: bisect_left(s, x))

    def count_leq(self, l: int, r: int, x: int) -> int:
        """Number of k in [l, r) with a[k] <= x."""
        return self._count(l, r, lambda s: bisect_right(s, x))

    def successor(self, l: int, r: int, x: int) -> Optional[int]:
        """Least a[k] >= x with k in [l, r), or None."""
        self._check_range(l, r)
        best: Optional[int] = None
        l += self.size
        r += self.size
        nodes: list[int] = []
        while l < r:
            if l & 1:
                nodes.append(l)
                l += 1
            if r & 1:
                r -= 1
                nodes.append(r)
            l //= 2
            r //= 2
        for v in nodes:
            s = self.lists[v]
            j = bisect_left(s, x)
            if j < len(s) and (best is None or s[j] < best):
                best = s[j]
        return best

    def kth_smallest(self, l: int, r: int, k: int) -> int:
        """The k-th smallest item of a[l:r], with k counted from 1."""
        self._check_range(l, r)
        if not 1 <= k <= r - l:
            raise ValueError(f"k = {k} out of range for {r - l} items")
        root = self.lists[1]
        lo, hi = 0, self.n - 1  # @line L9
        while lo < hi:  # @line L10
            m = (lo + hi) // 2  # @line L11
            if self.count_leq(l, r, root[m]) >= k:  # @line L12
                hi = m
            else:
                lo = m + 1
        return root[lo]  # @line L13

    def check_invariant(self) -> None:
        """Each node list is the sorted merge of its children's lists."""
        for v in range(self.size, 2 * self.size):
            assert len(self.lists[v]) == (1 if v - self.size < self.n else 0)
        for v in range(1, self.size):
            s = self.lists[v]
            assert s == sorted(self.lists[2 * v] + self.lists[2 * v + 1])
        assert len(self.lists[1]) == self.n