Spaces:
Running
Running
File size: 10,185 Bytes
b200bda |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 |
"""Priority queue class with updatable priorities.
"""
import heapq
__all__ = ["MappedQueue"]
class _HeapElement:
"""This proxy class separates the heap element from its priority.
The idea is that using a 2-tuple (priority, element) works
for sorting, but not for dict lookup because priorities are
often floating point values so round-off can mess up equality.
So, we need inequalities to look at the priority (for sorting)
and equality (and hash) to look at the element to enable
updates to the priority.
Unfortunately, this class can be tricky to work with if you forget that
`__lt__` compares the priority while `__eq__` compares the element.
In `greedy_modularity_communities()` the following code is
used to check that two _HeapElements differ in either element or priority:
if d_oldmax != row_max or d_oldmax.priority != row_max.priority:
If the priorities are the same, this implementation uses the element
as a tiebreaker. This provides compatibility with older systems that
use tuples to combine priority and elements.
"""
__slots__ = ["priority", "element", "_hash"]
def __init__(self, priority, element):
self.priority = priority
self.element = element
self._hash = hash(element)
def __lt__(self, other):
try:
other_priority = other.priority
except AttributeError:
return self.priority < other
# assume comparing to another _HeapElement
if self.priority == other_priority:
try:
return self.element < other.element
except TypeError as err:
raise TypeError(
"Consider using a tuple, with a priority value that can be compared."
)
return self.priority < other_priority
def __gt__(self, other):
try:
other_priority = other.priority
except AttributeError:
return self.priority > other
# assume comparing to another _HeapElement
if self.priority == other_priority:
try:
return self.element > other.element
except TypeError as err:
raise TypeError(
"Consider using a tuple, with a priority value that can be compared."
)
return self.priority > other_priority
def __eq__(self, other):
try:
return self.element == other.element
except AttributeError:
return self.element == other
def __hash__(self):
return self._hash
def __getitem__(self, indx):
return self.priority if indx == 0 else self.element[indx - 1]
def __iter__(self):
yield self.priority
try:
yield from self.element
except TypeError:
yield self.element
def __repr__(self):
return f"_HeapElement({self.priority}, {self.element})"
class MappedQueue:
"""The MappedQueue class implements a min-heap with removal and update-priority.
The min heap uses heapq as well as custom written _siftup and _siftdown
methods to allow the heap positions to be tracked by an additional dict
keyed by element to position. The smallest element can be popped in O(1) time,
new elements can be pushed in O(log n) time, and any element can be removed
or updated in O(log n) time. The queue cannot contain duplicate elements
and an attempt to push an element already in the queue will have no effect.
MappedQueue complements the heapq package from the python standard
library. While MappedQueue is designed for maximum compatibility with
heapq, it adds element removal, lookup, and priority update.
Parameters
----------
data : dict or iterable
Examples
--------
A `MappedQueue` can be created empty, or optionally, given a dictionary
of initial elements and priorities. The methods `push`, `pop`,
`remove`, and `update` operate on the queue.
>>> colors_nm = {'red':665, 'blue': 470, 'green': 550}
>>> q = MappedQueue(colors_nm)
>>> q.remove('red')
>>> q.update('green', 'violet', 400)
>>> q.push('indigo', 425)
True
>>> [q.pop().element for i in range(len(q.heap))]
['violet', 'indigo', 'blue']
A `MappedQueue` can also be initialized with a list or other iterable. The priority is assumed
to be the sort order of the items in the list.
>>> q = MappedQueue([916, 50, 4609, 493, 237])
>>> q.remove(493)
>>> q.update(237, 1117)
>>> [q.pop() for i in range(len(q.heap))]
[50, 916, 1117, 4609]
An exception is raised if the elements are not comparable.
>>> q = MappedQueue([100, 'a'])
Traceback (most recent call last):
...
TypeError: '<' not supported between instances of 'int' and 'str'
To avoid the exception, use a dictionary to assign priorities to the elements.
>>> q = MappedQueue({100: 0, 'a': 1 })
References
----------
.. [1] Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2001).
Introduction to algorithms second edition.
.. [2] Knuth, D. E. (1997). The art of computer programming (Vol. 3).
Pearson Education.
"""
def __init__(self, data=None):
"""Priority queue class with updatable priorities."""
if data is None:
self.heap = []
elif isinstance(data, dict):
self.heap = [_HeapElement(v, k) for k, v in data.items()]
else:
self.heap = list(data)
self.position = {}
self._heapify()
def _heapify(self):
"""Restore heap invariant and recalculate map."""
heapq.heapify(self.heap)
self.position = {elt: pos for pos, elt in enumerate(self.heap)}
if len(self.heap) != len(self.position):
raise AssertionError("Heap contains duplicate elements")
def __len__(self):
return len(self.heap)
def push(self, elt, priority=None):
"""Add an element to the queue."""
if priority is not None:
elt = _HeapElement(priority, elt)
# If element is already in queue, do nothing
if elt in self.position:
return False
# Add element to heap and dict
pos = len(self.heap)
self.heap.append(elt)
self.position[elt] = pos
# Restore invariant by sifting down
self._siftdown(0, pos)
return True
def pop(self):
"""Remove and return the smallest element in the queue."""
# Remove smallest element
elt = self.heap[0]
del self.position[elt]
# If elt is last item, remove and return
if len(self.heap) == 1:
self.heap.pop()
return elt
# Replace root with last element
last = self.heap.pop()
self.heap[0] = last
self.position[last] = 0
# Restore invariant by sifting up
self._siftup(0)
# Return smallest element
return elt
def update(self, elt, new, priority=None):
"""Replace an element in the queue with a new one."""
if priority is not None:
new = _HeapElement(priority, new)
# Replace
pos = self.position[elt]
self.heap[pos] = new
del self.position[elt]
self.position[new] = pos
# Restore invariant by sifting up
self._siftup(pos)
def remove(self, elt):
"""Remove an element from the queue."""
# Find and remove element
try:
pos = self.position[elt]
del self.position[elt]
except KeyError:
# Not in queue
raise
# If elt is last item, remove and return
if pos == len(self.heap) - 1:
self.heap.pop()
return
# Replace elt with last element
last = self.heap.pop()
self.heap[pos] = last
self.position[last] = pos
# Restore invariant by sifting up
self._siftup(pos)
def _siftup(self, pos):
"""Move smaller child up until hitting a leaf.
Built to mimic code for heapq._siftup
only updating position dict too.
"""
heap, position = self.heap, self.position
end_pos = len(heap)
startpos = pos
newitem = heap[pos]
# Shift up the smaller child until hitting a leaf
child_pos = (pos << 1) + 1 # start with leftmost child position
while child_pos < end_pos:
# Set child_pos to index of smaller child.
child = heap[child_pos]
right_pos = child_pos + 1
if right_pos < end_pos:
right = heap[right_pos]
if not child < right:
child = right
child_pos = right_pos
# Move the smaller child up.
heap[pos] = child
position[child] = pos
pos = child_pos
child_pos = (pos << 1) + 1
# pos is a leaf position. Put newitem there, and bubble it up
# to its final resting place (by sifting its parents down).
while pos > 0:
parent_pos = (pos - 1) >> 1
parent = heap[parent_pos]
if not newitem < parent:
break
heap[pos] = parent
position[parent] = pos
pos = parent_pos
heap[pos] = newitem
position[newitem] = pos
def _siftdown(self, start_pos, pos):
"""Restore invariant. keep swapping with parent until smaller.
Built to mimic code for heapq._siftdown
only updating position dict too.
"""
heap, position = self.heap, self.position
newitem = heap[pos]
# Follow the path to the root, moving parents down until finding a place
# newitem fits.
while pos > start_pos:
parent_pos = (pos - 1) >> 1
parent = heap[parent_pos]
if not newitem < parent:
break
heap[pos] = parent
position[parent] = pos
pos = parent_pos
heap[pos] = newitem
position[newitem] = pos
|