Mercurial > repos > shellac > sam_consensus_v3
comparison env/lib/python3.9/site-packages/networkx/utils/mapped_queue.py @ 0:4f3585e2f14b draft default tip
"planemo upload commit 60cee0fc7c0cda8592644e1aad72851dec82c959"
author | shellac |
---|---|
date | Mon, 22 Mar 2021 18:12:50 +0000 |
parents | |
children |
comparison
equal
deleted
inserted
replaced
-1:000000000000 | 0:4f3585e2f14b |
---|---|
1 """Priority queue class with updatable priorities. | |
2 """ | |
3 | |
4 import heapq | |
5 | |
6 __all__ = ["MappedQueue"] | |
7 | |
8 | |
9 class MappedQueue: | |
10 """The MappedQueue class implements an efficient minimum heap. The | |
11 smallest element can be popped in O(1) time, new elements can be pushed | |
12 in O(log n) time, and any element can be removed or updated in O(log n) | |
13 time. The queue cannot contain duplicate elements and an attempt to push an | |
14 element already in the queue will have no effect. | |
15 | |
16 MappedQueue complements the heapq package from the python standard | |
17 library. While MappedQueue is designed for maximum compatibility with | |
18 heapq, it has slightly different functionality. | |
19 | |
20 Examples | |
21 -------- | |
22 | |
23 A `MappedQueue` can be created empty or optionally given an array of | |
24 initial elements. Calling `push()` will add an element and calling `pop()` | |
25 will remove and return the smallest element. | |
26 | |
27 >>> q = MappedQueue([916, 50, 4609, 493, 237]) | |
28 >>> q.push(1310) | |
29 True | |
30 >>> x = [q.pop() for i in range(len(q.h))] | |
31 >>> x | |
32 [50, 237, 493, 916, 1310, 4609] | |
33 | |
34 Elements can also be updated or removed from anywhere in the queue. | |
35 | |
36 >>> q = MappedQueue([916, 50, 4609, 493, 237]) | |
37 >>> q.remove(493) | |
38 >>> q.update(237, 1117) | |
39 >>> x = [q.pop() for i in range(len(q.h))] | |
40 >>> x | |
41 [50, 916, 1117, 4609] | |
42 | |
43 References | |
44 ---------- | |
45 .. [1] Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2001). | |
46 Introduction to algorithms second edition. | |
47 .. [2] Knuth, D. E. (1997). The art of computer programming (Vol. 3). | |
48 Pearson Education. | |
49 """ | |
50 | |
51 def __init__(self, data=[]): | |
52 """Priority queue class with updatable priorities. | |
53 """ | |
54 self.h = list(data) | |
55 self.d = dict() | |
56 self._heapify() | |
57 | |
58 def __len__(self): | |
59 return len(self.h) | |
60 | |
61 def _heapify(self): | |
62 """Restore heap invariant and recalculate map.""" | |
63 heapq.heapify(self.h) | |
64 self.d = {elt: pos for pos, elt in enumerate(self.h)} | |
65 if len(self.h) != len(self.d): | |
66 raise AssertionError("Heap contains duplicate elements") | |
67 | |
68 def push(self, elt): | |
69 """Add an element to the queue.""" | |
70 # If element is already in queue, do nothing | |
71 if elt in self.d: | |
72 return False | |
73 # Add element to heap and dict | |
74 pos = len(self.h) | |
75 self.h.append(elt) | |
76 self.d[elt] = pos | |
77 # Restore invariant by sifting down | |
78 self._siftdown(pos) | |
79 return True | |
80 | |
81 def pop(self): | |
82 """Remove and return the smallest element in the queue.""" | |
83 # Remove smallest element | |
84 elt = self.h[0] | |
85 del self.d[elt] | |
86 # If elt is last item, remove and return | |
87 if len(self.h) == 1: | |
88 self.h.pop() | |
89 return elt | |
90 # Replace root with last element | |
91 last = self.h.pop() | |
92 self.h[0] = last | |
93 self.d[last] = 0 | |
94 # Restore invariant by sifting up, then down | |
95 pos = self._siftup(0) | |
96 self._siftdown(pos) | |
97 # Return smallest element | |
98 return elt | |
99 | |
100 def update(self, elt, new): | |
101 """Replace an element in the queue with a new one.""" | |
102 # Replace | |
103 pos = self.d[elt] | |
104 self.h[pos] = new | |
105 del self.d[elt] | |
106 self.d[new] = pos | |
107 # Restore invariant by sifting up, then down | |
108 pos = self._siftup(pos) | |
109 self._siftdown(pos) | |
110 | |
111 def remove(self, elt): | |
112 """Remove an element from the queue.""" | |
113 # Find and remove element | |
114 try: | |
115 pos = self.d[elt] | |
116 del self.d[elt] | |
117 except KeyError: | |
118 # Not in queue | |
119 raise | |
120 # If elt is last item, remove and return | |
121 if pos == len(self.h) - 1: | |
122 self.h.pop() | |
123 return | |
124 # Replace elt with last element | |
125 last = self.h.pop() | |
126 self.h[pos] = last | |
127 self.d[last] = pos | |
128 # Restore invariant by sifting up, then down | |
129 pos = self._siftup(pos) | |
130 self._siftdown(pos) | |
131 | |
132 def _siftup(self, pos): | |
133 """Move element at pos down to a leaf by repeatedly moving the smaller | |
134 child up.""" | |
135 h, d = self.h, self.d | |
136 elt = h[pos] | |
137 # Continue until element is in a leaf | |
138 end_pos = len(h) | |
139 left_pos = (pos << 1) + 1 | |
140 while left_pos < end_pos: | |
141 # Left child is guaranteed to exist by loop predicate | |
142 left = h[left_pos] | |
143 try: | |
144 right_pos = left_pos + 1 | |
145 right = h[right_pos] | |
146 # Out-of-place, swap with left unless right is smaller | |
147 if right < left: | |
148 h[pos], h[right_pos] = right, elt | |
149 pos, right_pos = right_pos, pos | |
150 d[elt], d[right] = pos, right_pos | |
151 else: | |
152 h[pos], h[left_pos] = left, elt | |
153 pos, left_pos = left_pos, pos | |
154 d[elt], d[left] = pos, left_pos | |
155 except IndexError: | |
156 # Left leaf is the end of the heap, swap | |
157 h[pos], h[left_pos] = left, elt | |
158 pos, left_pos = left_pos, pos | |
159 d[elt], d[left] = pos, left_pos | |
160 # Update left_pos | |
161 left_pos = (pos << 1) + 1 | |
162 return pos | |
163 | |
164 def _siftdown(self, pos): | |
165 """Restore invariant by repeatedly replacing out-of-place element with | |
166 its parent.""" | |
167 h, d = self.h, self.d | |
168 elt = h[pos] | |
169 # Continue until element is at root | |
170 while pos > 0: | |
171 parent_pos = (pos - 1) >> 1 | |
172 parent = h[parent_pos] | |
173 if parent > elt: | |
174 # Swap out-of-place element with parent | |
175 h[parent_pos], h[pos] = elt, parent | |
176 parent_pos, pos = pos, parent_pos | |
177 d[elt] = pos | |
178 d[parent] = parent_pos | |
179 else: | |
180 # Invariant is satisfied | |
181 break | |
182 return pos |