Source code for titan_pylib.data_structures.heap.radix_heap

 1from typing import Generic, TypeVar
 2
 3T = TypeVar("T")
 4
 5

[docs]
 6class RadixHeap(Generic[T]):
 7
 8    def __init__(self, u: int):
 9        self.u = u
10        self.log = u.bit_length()
11        self.lim = (1 << self.log) - 1
12        self.last = 0
13        self._len = 0
14        self.data: list[list[tuple[int, T]]] = [[] for _ in range(self.log)]
15

[docs]
16    def push(self, key: int, val: T) -> None:
17        assert key <= self.lim
18        self._len += 1
19        self.data[(key ^ self.last).bit_length()].append((key, val))

20

[docs]
21    def pop_min(self) -> tuple[int, T]:
22        self._len -= 1
23        data = self.data
24        if data[0]:
25            return data[0].pop()
26        for d in data:
27            if not d:
28                continue
29            last = min(d)[0]
30            for elm in d:
31                data[(elm[0] ^ last).bit_length()].append(elm)
32            d.clear()
33            self.last = last
34            break
35        return data[0].pop()

36
37    def __len__(self):
38        return self._len
39
40    def __bool__(self):
41        return self._len > 0
42
43    def __str__(self):
44        a = []
45        for d in self.data:
46            a.extend(d)
47        return str(a)