Indexed priority queue in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/indexed-priority-queue/. The authors are listed in prov/QQ610A.jsonl.

from typing import Any, Generic, Optional, Protocol, TypeVar


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


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


class IndexedMinPQ(Generic[T]):
    """A binary min-heap of item indices 0 .. m - 1, ordered by their keys.

    pq is the heap, a 0-based list of indices; qp[i] is the position of the
    index i in pq, or -1 if i is not in the queue; key[i] is the key of i.
    So pq[qp[i]] == i for each i in the queue. insert, extract_min,
    decrease_key, change_key and delete take O(log n) time; contains and
    key_of take O(1) time. Keys need only "<".
    """

    def __init__(self, m: int) -> None:
        if m < 0:
            raise ValueError("capacity must be at least 0")
        self.m = m
        self.pq: list[int] = []
        self.qp: list[int] = [-1] * m
        self.key: list[Optional[T]] = [None] * m

    def __len__(self) -> int:
        return len(self.pq)

    def _check_index(self, i: int) -> None:
        if not 0 <= i < self.m:
            raise IndexError("index out of range")

    def contains(self, i: int) -> bool:
        self._check_index(i)
        return self.qp[i] != -1

    def key_of(self, i: int) -> T:
        if not self.contains(i):
            raise KeyError("index not in the queue")
        k = self.key[i]
        assert k is not None
        return k

    def _less(self, a: int, b: int) -> bool:
        ka, kb = self.key[self.pq[a]], self.key[self.pq[b]]
        assert ka is not None and kb is not None
        return ka < kb

    def _exch(self, a: int, b: int) -> None:
        pq = self.pq
        pq[a], pq[b] = pq[b], pq[a]  # @line L13
        self.qp[pq[a]] = a
        self.qp[pq[b]] = b

    def _swim(self, j: int) -> None:
        while j > 0 and self._less(j, (j - 1) // 2):  # @line L14
            self._exch(j, (j - 1) // 2)  # @line L15
            j = (j - 1) // 2

    def _sink(self, j: int) -> None:
        n = len(self.pq)
        while 2 * j + 1 < n:  # @line L16
            c = 2 * j + 1  # @line L17
            if c + 1 < n and self._less(c + 1, c):
                c += 1
            if not self._less(c, j):  # @line L18
                return
            self._exch(j, c)  # @line L19
            j = c

    def insert(self, i: int, k: T) -> None:
        if self.contains(i):  # @line L1
            raise KeyError("index already in the queue")
        self.qp[i] = len(self.pq)  # @line L2
        self.pq.append(i)
        self.key[i] = k
        self._swim(self.qp[i])  # @line L3

    def min_index(self) -> int:
        if not self.pq:
            raise IndexError("min_index from an empty queue")
        return self.pq[0]

    def extract_min(self) -> tuple[int, T]:
        if not self.pq:  # @line L4
            raise IndexError("extract_min from an empty queue")
        i = self.pq[0]  # @line L5
        k = self.key_of(i)
        self._exch(0, len(self.pq) - 1)
        self.pq.pop()
        self._sink(0)  # @line L6
        self.qp[i] = -1
        self.key[i] = None
        return i, k  # @line L7

    def decrease_key(self, i: int, k: T) -> None:
        if self.key_of(i) < k:  # @line L8
            raise ValueError("new key is larger than the current key")
        self.key[i] = k  # @line L9
        self._swim(self.qp[i])

    def change_key(self, i: int, k: T) -> None:
        old = self.key_of(i)
        self.key[i] = k
        if k < old:
            self._swim(self.qp[i])
        else:
            self._sink(self.qp[i])

    def delete(self, i: int) -> T:
        k = self.key_of(i)
        j = self.qp[i]  # @line L10
        self._exch(j, len(self.pq) - 1)
        self.pq.pop()
        if j < len(self.pq):  # @line L11
            self._swim(j)
            self._sink(j)
        self.qp[i] = -1  # @line L12
        self.key[i] = None
        return k

    def check_invariant(self) -> None:
        n = len(self.pq)
        for j, i in enumerate(self.pq):
            assert self.qp[i] == j
            assert self.key[i] is not None
            if j > 0:
                assert not self._less(j, (j - 1) // 2)
        assert sum(1 for p in self.qp if p != -1) == n
        for i in range(self.m):
            if self.qp[i] == -1:
                assert self.key[i] is None