class HeapDict:
__slots__ = ('_dict', '_heap', '__getitem__', '__delitem__', '__len__',
'__iter__', 'pop') # fmt: skip
def __init__(self, dict=None, /, **kwargs):
self._dict = {}
if dict is not None:
self._dict.update(dict)
if kwargs:
self._dict.update(kwargs)
self.__getitem__ = self._dict.__getitem__
self.__delitem__ = self._dict.__delitem__
self.__len__ = self._dict.__len__
self.__iter__ = self._dict.__iter__
self.pop = self._dict.pop
self._heap = [(v, k) for k, v in self._dict.items()]
heapq.heapify(self._heap)
def __setitem__(self, key, comp_value):
self._dict[key] = comp_value
heapq.heappush(self._heap, (comp_value, key))
def extract_min(self):
val, key = heapq.heappop(self._heap)
while self._dict.get(key) != val:
val, key = heapq.heappop(self._heap)
del self._dict[key]
return key, val
def find_min(self):
val, key = self._heap[0]
while self._dict.get(key) != val:
heapq.heappop(self._heap)
val, key = self._heap[0]
return key, val
class HeapDict(collections.UserDict):
__slots__ = ('_heap',)
def __init__(self, mapping=None, /, **kwargs):
super().__init__()
if dict is not None:
self.data.update(mapping)
if kwargs:
self.data.update(kwargs)
self._heap = [(v, k) for k, v in self.data.items()]
heapq.heapify(self._heap)
def __setitem__(self, key, comp_value):
self.data[key] = comp_value
heapq.heappush(self._heap, (comp_value, key))
...
def diskstra(wgraph, source):
distances = [INF] * len(wgraph)
distances[source] = 0
heap = [(0, source)]
while heap:
dist_u, u = heapq.heappop(heap)
if dist_u != distances[u]:
continue
for v, dist_uv in wgraph[u].items():
if (dist_v := dist_u + dist_uv) < distances[v]:
distances[v] = dist_v
heapq.heappush(heap, (dist_v, v))
return distances
def dijkstra(wgraph, source):
distances = [INF] * len(wgraph)
distances[source] = 0
Q = PriorityDict({source: 0})
while Q:
u, dist_u = Q.pop_min()
for v, dist_uv in wgraph[u].items():
if (dist_v := dist_u + dist_uv) < distances[v]:
distances[v] = dist_v
Q[v] = dist_v
return distances