Calendar queue in Python

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

import bisect
import heapq
import math
from typing import Generic, Optional, TypeVar

E = TypeVar("E")


class CalendarQueue(Generic[E]):
    """Brown's calendar queue: a min-priority queue of events with real times.

    There are N buckets of width w. An event with time t has the virtual
    bucket vb(t) = floor(t / w) and lives in the bucket vb(t) mod N, which is
    kept sorted by (time, arrival). V is a virtual bucket with vb(t) >= V for
    every event in the queue. extract_min scans at most one year of N
    buckets from V, then falls back to a direct search. With resize=True,
    N doubles when the size exceeds 2N and halves when it falls below N/2,
    and w is estimated from the gaps between the first events.
    """

    def __init__(self, nbuckets: int = 2, width: float = 1.0, resize: bool = True) -> None:
        if nbuckets < 1 or not width > 0:
            raise ValueError("nbuckets must be at least 1 and width positive")
        self.n_buckets = nbuckets
        self.width = width
        self.resize_on = resize
        self.buckets: list[list[tuple[float, int, E]]] = [[] for _ in range(nbuckets)]
        self.size = 0
        self.v = 0
        self.seq = 0
        self.resizes = 0

    def __len__(self) -> int:
        return self.size

    def _vb(self, t: float) -> int:
        return math.floor(t / self.width)

    def insert(self, e: E, t: float) -> None:
        vb = self._vb(t)
        i = vb % self.n_buckets  # @line L1
        bisect.insort(self.buckets[i], (t, self.seq, e))  # @line L2
        self.seq += 1
        if self.size == 0 or vb < self.v:  # @line L3
            self.v = vb
        self.size += 1
        if self.resize_on and self.size > 2 * self.n_buckets:  # @line L4
            self._resize(2 * self.n_buckets)

    def _take(self, i: int) -> tuple[E, float]:
        t, _, e = self.buckets[i].pop(0)
        self.size -= 1
        if self.resize_on and self.n_buckets > 2 and self.size < self.n_buckets // 2:
            self._resize(self.n_buckets // 2)
        return e, t

    def peek_min(self) -> tuple[E, float]:
        if self.size == 0:
            raise IndexError("peek_min from an empty queue")
        best = min(b[0] for b in self.buckets if b)
        return best[2], best[0]

    def extract_min(self) -> tuple[E, float]:
        if self.size == 0:  # @line L5
            raise IndexError("extract_min from an empty queue")
        n = self.n_buckets
        for _ in range(n):  # @line L6
            b = self.buckets[self.v % n]
            if b and self._vb(b[0][0]) == self.v:  # @line L7
                return self._take(self.v % n)  # @line L8 @line L9
            self.v += 1  # @line L10
        best: Optional[tuple[float, int, E]] = None
        best_i = -1
        for i, b in enumerate(self.buckets):  # @line L11
            if b and (best is None or b[0] < best):
                best, best_i = b[0], i
        assert best is not None
        self.v = self._vb(best[0])  # @line L12
        return self._take(best_i)

    def _new_width(self, events: list[tuple[float, int, E]]) -> float:
        front = heapq.nsmallest(25, events)  # @line L14
        gaps = [b[0] - a[0] for a, b in zip(front, front[1:])]
        if not gaps:
            return self.width
        avg = sum(gaps) / len(gaps)
        kept = [g for g in gaps if g <= 2 * avg]
        w = 3 * sum(kept) / len(kept) if kept else 0.0  # @line L15
        return w if w > 0 else self.width

    def _resize(self, new_n: int) -> None:
        events = [x for b in self.buckets for x in b]  # @line L13
        self.width = self._new_width(events)
        self.n_buckets = new_n
        self.buckets = [[] for _ in range(new_n)]
        for x in events:
            self.buckets[self._vb(x[0]) % new_n].append(x)
        for b in self.buckets:
            b.sort()
        self.v = min((self._vb(x[0]) for x in events), default=0)
        self.resizes += 1

    def check_invariant(self) -> None:
        count = 0
        for i, b in enumerate(self.buckets):
            assert b == sorted(b)
            for t, _, _ in b:
                vb = self._vb(t)
                assert vb % self.n_buckets == i
                assert vb >= self.v
                count += 1
        assert count == self.size