Radix heap in Python
Monotone priority queue for integer keys using buckets indexed by the highest differing bit.
- File
radix_heap.py, 82 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/radix-heap/. The authors are listed in prov/FZSDPB.jsonl.
from typing import Generic, TypeVar
V = TypeVar("V")
class RadixHeap(Generic[V]):
"""A monotone min-priority queue on nonnegative integer keys.
Bucket 0 holds the items whose key equals last, the last extracted key.
Bucket i >= 1 holds the items whose key differs from last first in bit
i - 1, counting from bit 0. Every key inserted must be at least last.
"""
def __init__(self) -> None:
self.last = 0
self.buckets: list[list[tuple[int, V | None]]] = [[]]
self.n = 0
self.moves = 0
def __len__(self) -> int:
return self.n
def is_empty(self) -> bool:
"""Return True if the heap holds no item."""
return self.n == 0
def _bucket(self, key: int) -> int:
return (key ^ self.last).bit_length()
def _put(self, key: int, value: V | None) -> None:
i = self._bucket(key)
while len(self.buckets) <= i:
self.buckets.append([])
self.buckets[i].append((key, value)) # @line L2
def insert(self, key: int, value: V | None = None) -> None:
"""Add an item; key must be an integer at least the last extracted key."""
if key < self.last: # @line L1
raise ValueError("radix heap is monotone: key below the last extracted key")
self._put(key, value)
self.n += 1
def find_min(self) -> int:
"""Return a smallest key without removing it."""
if self.n == 0:
raise IndexError("find_min on an empty heap")
self._refill()
return self.last
def _refill(self) -> None:
if self.buckets[0]: # @line L3
return
i = 1 # @line L4
while not self.buckets[i]:
i += 1
items = self.buckets[i]
self.last = min(k for k, _ in items) # @line L5
self.buckets[i] = [] # @line L6
for k, v in items: # @line L7
self._put(k, v)
self.moves += 1
def extract_min(self) -> tuple[int, V | None]:
"""Remove and return an item (key, value) with the smallest key."""
if self.n == 0:
raise IndexError("extract_min on an empty heap")
self._refill()
self.n -= 1
return self.buckets[0].pop() # @line L8
def check_invariant(self) -> None:
"""Check that every item is in the bucket given by its key and last."""
count = 0
for i, b in enumerate(self.buckets):
for k, _ in b:
assert k >= self.last, "monotone"
assert self._bucket(k) == i, "bucket"
count += 1
assert count == self.n, "count"
test_radix_heap.py 116 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/radix-heap/. The authors are listed in prov/FZSDPB.jsonl.
import heapq
import random
import pytest
from radix_heap import RadixHeap
def contents(h: RadixHeap[None]) -> dict[int, list[int]]:
return {i: sorted(k for k, _ in b) for i, b in enumerate(h.buckets) if b}
def test_example() -> None:
h: RadixHeap[None] = RadixHeap()
for k in [5, 1, 7, 6, 13]:
h.insert(k)
h.check_invariant()
assert contents(h) == {1: [1], 3: [5, 6, 7], 4: [13]}
assert h.extract_min()[0] == 1
assert contents(h) == {3: [5, 6, 7], 4: [13]}
assert h.extract_min()[0] == 5
assert contents(h) == {2: [6, 7], 4: [13]}
h.insert(8)
assert contents(h) == {2: [6, 7], 4: [8, 13]}
assert h.extract_min()[0] == 6
assert contents(h) == {1: [7], 4: [8, 13]}
assert h.extract_min()[0] == 7
assert h.extract_min()[0] == 8
assert contents(h) == {3: [13]}
assert h.extract_min()[0] == 13
assert h.is_empty()
assert h.moves == 10
h.check_invariant()
def test_edge_cases() -> None:
h: RadixHeap[str] = RadixHeap()
with pytest.raises(IndexError):
h.extract_min()
with pytest.raises(IndexError):
h.find_min()
h.insert(0, "a")
h.insert(0, "b")
assert h.find_min() == 0
assert {h.extract_min()[1], h.extract_min()[1]} == {"a", "b"}
h.insert(10, "c")
assert h.extract_min() == (10, "c")
with pytest.raises(ValueError):
h.insert(9, "d")
big = 2**70 + 3
h.insert(big, "e")
h.insert(10, "f")
assert h.extract_min() == (10, "f")
assert h.extract_min() == (big, "e")
def test_random_monotone_against_heapq() -> None:
rng = random.Random(1990)
for _ in range(600):
h: RadixHeap[int] = RadixHeap()
ref: list[int] = []
last = 0
inserts = 0
top = 1 << rng.randrange(1, 20)
for _ in range(rng.randrange(1, 80)):
if rng.random() < 0.6 or not ref:
k = last + rng.randrange(top)
h.insert(k, k)
heapq.heappush(ref, k)
inserts += 1
else:
k, v = h.extract_min()
assert k == v == heapq.heappop(ref)
last = k
h.check_invariant()
assert len(h) == len(ref)
bits = max(1, (last + top).bit_length())
assert h.moves <= inserts * bits
def test_dijkstra() -> None:
rng = random.Random(42)
for _ in range(100):
n = rng.randrange(2, 40)
adj: list[list[tuple[int, int]]] = [[] for _ in range(n)]
for _ in range(rng.randrange(0, 4 * n)):
u, v = rng.randrange(n), rng.randrange(n)
adj[u].append((v, rng.randrange(0, 100)))
def run(use_radix: bool) -> list[float]:
dist = [float("inf")] * n
dist[0] = 0
h: RadixHeap[int] = RadixHeap()
q: list[tuple[int, int]] = []
if use_radix:
h.insert(0, 0)
else:
q.append((0, 0))
while (len(h) if use_radix else len(q)):
d, u = h.extract_min() if use_radix else heapq.heappop(q)
assert u is not None
if d > dist[u]:
continue
for v, w in adj[u]:
if d + w < dist[v]:
dist[v] = d + w
if use_radix:
h.insert(d + w, v)
else:
heapq.heappush(q, (d + w, v))
return dist
assert run(True) == run(False)
What the next level needs
- the tests did not run in this build