1"""
2Min-heaps.
3"""
4
5from heapq import heappop, heappush
6from itertools import count
7
8import networkx as nx
9
10__all__ = ["MinHeap", "PairingHeap", "BinaryHeap"]
11
12
13class MinHeap:
14 """Base class for min-heaps.
15
16 A MinHeap stores a collection of key-value pairs ordered by their values.
17 It supports querying the minimum pair, inserting a new pair, decreasing the
18 value in an existing pair and deleting the minimum pair.
19 """
20
21 class _Item:
22 """Used by subclassess to represent a key-value pair."""
23
24 __slots__ = ("key", "value")
25
26 def __init__(self, key, value):
27 self.key = key
28 self.value = value
29
30 def __repr__(self):
31 return repr((self.key, self.value))
32
33 def __init__(self):
34 """Initialize a new min-heap."""
35 self._dict = {}
36
37 def min(self):
38 """Query the minimum key-value pair.
39
40 Returns
41 -------
42 key, value : tuple
43 The key-value pair with the minimum value in the heap.
44
45 Raises
46 ------
47 NetworkXError
48 If the heap is empty.
49 """
50 raise NotImplementedError
51
52 def pop(self):
53 """Delete the minimum pair in the heap.
54
55 Returns
56 -------
57 key, value : tuple
58 The key-value pair with the minimum value in the heap.
59
60 Raises
61 ------
62 NetworkXError
63 If the heap is empty.
64 """
65 raise NotImplementedError
66
67 def get(self, key, default=None):
68 """Returns the value associated with a key.
69
70 Parameters
71 ----------
72 key : hashable object
73 The key to be looked up.
74
75 default : object
76 Default value to return if the key is not present in the heap.
77 Default value: None.
78
79 Returns
80 -------
81 value : object.
82 The value associated with the key.
83 """
84 raise NotImplementedError
85
86 def insert(self, key, value, allow_increase=False):
87 """Insert a new key-value pair or modify the value in an existing
88 pair.
89
90 Parameters
91 ----------
92 key : hashable object
93 The key.
94
95 value : object comparable with existing values.
96 The value.
97
98 allow_increase : bool
99 Whether the value is allowed to increase. If False, attempts to
100 increase an existing value have no effect. Default value: False.
101
102 Returns
103 -------
104 decreased : bool
105 True if a pair is inserted or the existing value is decreased.
106 """
107 raise NotImplementedError
108
109 def __nonzero__(self):
110 """Returns whether the heap if empty."""
111 return bool(self._dict)
112
113 def __bool__(self):
114 """Returns whether the heap if empty."""
115 return bool(self._dict)
116
117 def __len__(self):
118 """Returns the number of key-value pairs in the heap."""
119 return len(self._dict)
120
121 def __contains__(self, key):
122 """Returns whether a key exists in the heap.
123
124 Parameters
125 ----------
126 key : any hashable object.
127 The key to be looked up.
128 """
129 return key in self._dict
130
131
132class PairingHeap(MinHeap):
133 """A pairing heap."""
134
135 class _Node(MinHeap._Item):
136 """A node in a pairing heap.
137
138 A tree in a pairing heap is stored using the left-child, right-sibling
139 representation.
140 """
141
142 __slots__ = ("left", "next", "prev", "parent")
143
144 def __init__(self, key, value):
145 super().__init__(key, value)
146 # The leftmost child.
147 self.left = None
148 # The next sibling.
149 self.next = None
150 # The previous sibling.
151 self.prev = None
152 # The parent.
153 self.parent = None
154
155 def __init__(self):
156 """Initialize a pairing heap."""
157 super().__init__()
158 self._root = None
159
160 def min(self):
161 if self._root is None:
162 raise nx.NetworkXError("heap is empty.")
163 return (self._root.key, self._root.value)
164
165 def pop(self):
166 if self._root is None:
167 raise nx.NetworkXError("heap is empty.")
168 min_node = self._root
169 self._root = self._merge_children(self._root)
170 del self._dict[min_node.key]
171 return (min_node.key, min_node.value)
172
173 def get(self, key, default=None):
174 node = self._dict.get(key)
175 return node.value if node is not None else default
176
177 def insert(self, key, value, allow_increase=False):
178 node = self._dict.get(key)
179 root = self._root
180 if node is not None:
181 if value < node.value:
182 node.value = value
183 if node is not root and value < node.parent.value:
184 self._cut(node)
185 self._root = self._link(root, node)
186 return True
187 elif allow_increase and value > node.value:
188 node.value = value
189 child = self._merge_children(node)
190 # Nonstandard step: Link the merged subtree with the root. See
191 # below for the standard step.
192 if child is not None:
193 self._root = self._link(self._root, child)
194 # Standard step: Perform a decrease followed by a pop as if the
195 # value were the smallest in the heap. Then insert the new
196 # value into the heap.
197 # if node is not root:
198 # self._cut(node)
199 # if child is not None:
200 # root = self._link(root, child)
201 # self._root = self._link(root, node)
202 # else:
203 # self._root = (self._link(node, child)
204 # if child is not None else node)
205 return False
206 else:
207 # Insert a new key.
208 node = self._Node(key, value)
209 self._dict[key] = node
210 self._root = self._link(root, node) if root is not None else node
211 return True
212
213 def _link(self, root, other):
214 """Link two nodes, making the one with the smaller value the parent of
215 the other.
216 """
217 if other.value < root.value:
218 root, other = other, root
219 next = root.left
220 other.next = next
221 if next is not None:
222 next.prev = other
223 other.prev = None
224 root.left = other
225 other.parent = root
226 return root
227
228 def _merge_children(self, root):
229 """Merge the subtrees of the root using the standard two-pass method.
230 The resulting subtree is detached from the root.
231 """
232 node = root.left
233 root.left = None
234 if node is not None:
235 link = self._link
236 # Pass 1: Merge pairs of consecutive subtrees from left to right.
237 # At the end of the pass, only the prev pointers of the resulting
238 # subtrees have meaningful values. The other pointers will be fixed
239 # in pass 2.
240 prev = None
241 while True:
242 next = node.next
243 if next is None:
244 node.prev = prev
245 break
246 next_next = next.next
247 node = link(node, next)
248 node.prev = prev
249 prev = node
250 if next_next is None:
251 break
252 node = next_next
253 # Pass 2: Successively merge the subtrees produced by pass 1 from
254 # right to left with the rightmost one.
255 prev = node.prev
256 while prev is not None:
257 prev_prev = prev.prev
258 node = link(prev, node)
259 prev = prev_prev
260 # Now node can become the new root. Its has no parent nor siblings.
261 node.prev = None
262 node.next = None
263 node.parent = None
264 return node
265
266 def _cut(self, node):
267 """Cut a node from its parent."""
268 prev = node.prev
269 next = node.next
270 if prev is not None:
271 prev.next = next
272 else:
273 node.parent.left = next
274 node.prev = None
275 if next is not None:
276 next.prev = prev
277 node.next = None
278 node.parent = None
279
280
281class BinaryHeap(MinHeap):
282 """A binary heap."""
283
284 def __init__(self):
285 """Initialize a binary heap."""
286 super().__init__()
287 self._heap = []
288 self._count = count()
289
290 def min(self):
291 dict = self._dict
292 if not dict:
293 raise nx.NetworkXError("heap is empty")
294 heap = self._heap
295 # Repeatedly remove stale key-value pairs until a up-to-date one is
296 # met.
297 while True:
298 value, _, key = heap[0]
299 if key in dict and value == dict[key]:
300 break
301 heappop(heap)
302 return (key, value)
303
304 def pop(self):
305 dict = self._dict
306 if not dict:
307 raise nx.NetworkXError("heap is empty")
308 heap = self._heap
309 # Repeatedly remove stale key-value pairs until a up-to-date one is
310 # met.
311 while True:
312 value, _, key = heap[0]
313 heappop(heap)
314 if key in dict and value == dict[key]:
315 break
316 del dict[key]
317 return (key, value)
318
319 def get(self, key, default=None):
320 return self._dict.get(key, default)
321
322 def insert(self, key, value, allow_increase=False):
323 dict = self._dict
324 if key in dict:
325 old_value = dict[key]
326 if value < old_value or (allow_increase and value > old_value):
327 # Since there is no way to efficiently obtain the location of a
328 # key-value pair in the heap, insert a new pair even if ones
329 # with the same key may already be present. Deem the old ones
330 # as stale and skip them when the minimum pair is queried.
331 dict[key] = value
332 heappush(self._heap, (value, next(self._count), key))
333 return value < old_value
334 return False
335 else:
336 dict[key] = value
337 heappush(self._heap, (value, next(self._count), key))
338 return True