ioftools / networkxMiCe / networkxmaster / networkx / utils / mapped_queue.py @ 5cef0f13
History  View  Annotate  Download (6.15 KB)
1 
# * coding: utf8 *


2 
#

3 
# priorityq: An objectoriented priority queue with updatable priorities.

4 
#

5 
# Copyright 2018 Edward L. Platt

6 
#

7 
# This file is part of NetworkX

8 
#

9 
# NetworkX is distributed under a BSD license; see LICENSE.txt for more

10 
# information.

11 
#

12 
# Authors:

13 
# Edward L. Platt <ed@elplatt.com>

14 
#

15 
"""Priority queue class with updatable priorities.

16 
"""

17  
18 
import heapq 
19  
20 
__all__ = ['MappedQueue']

21  
22  
23 
class MappedQueue(object): 
24 
"""The MappedQueue class implements an efficient minimum heap. The

25 
smallest element can be popped in O(1) time, new elements can be pushed

26 
in O(log n) time, and any element can be removed or updated in O(log n)

27 
time. The queue cannot contain duplicate elements and an attempt to push an

28 
element already in the queue will have no effect.

29 

30 
MappedQueue complements the heapq package from the python standard

31 
library. While MappedQueue is designed for maximum compatibility with

32 
heapq, it has slightly different functionality.

33 

34 
Examples

35 


36 

37 
A `MappedQueue` can be created empty or optionally given an array of

38 
initial elements. Calling `push()` will add an element and calling `pop()`

39 
will remove and return the smallest element.

40 

41 
>>> q = MappedQueue([916, 50, 4609, 493, 237])

42 
>>> q.push(1310)

43 
True

44 
>>> x = [q.pop() for i in range(len(q.h))]

45 
>>> x

46 
[50, 237, 493, 916, 1310, 4609]

47 

48 
Elements can also be updated or removed from anywhere in the queue.

49 

50 
>>> q = MappedQueue([916, 50, 4609, 493, 237])

51 
>>> q.remove(493)

52 
>>> q.update(237, 1117)

53 
>>> x = [q.pop() for i in range(len(q.h))]

54 
>>> x

55 
[50, 916, 1117, 4609]

56 

57 
References

58 


59 
.. [1] Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2001).

60 
Introduction to algorithms second edition.

61 
.. [2] Knuth, D. E. (1997). The art of computer programming (Vol. 3).

62 
Pearson Education.

63 
"""

64  
65 
def __init__(self, data=[]): 
66 
"""Priority queue class with updatable priorities.

67 
"""

68 
self.h = list(data) 
69 
self.d = dict() 
70 
self._heapify()

71  
72 
def __len__(self): 
73 
return len(self.h) 
74  
75 
def _heapify(self): 
76 
"""Restore heap invariant and recalculate map."""

77 
heapq.heapify(self.h)

78 
self.d = dict([(elt, pos) for pos, elt in enumerate(self.h)]) 
79 
if len(self.h) != len(self.d): 
80 
raise AssertionError("Heap contains duplicate elements") 
81  
82 
def push(self, elt): 
83 
"""Add an element to the queue."""

84 
# If element is already in queue, do nothing

85 
if elt in self.d: 
86 
return False 
87 
# Add element to heap and dict

88 
pos = len(self.h) 
89 
self.h.append(elt)

90 
self.d[elt] = pos

91 
# Restore invariant by sifting down

92 
self._siftdown(pos)

93 
return True 
94  
95 
def pop(self): 
96 
"""Remove and return the smallest element in the queue."""

97 
# Remove smallest element

98 
elt = self.h[0] 
99 
del self.d[elt] 
100 
# If elt is last item, remove and return

101 
if len(self.h) == 1: 
102 
self.h.pop()

103 
return elt

104 
# Replace root with last element

105 
last = self.h.pop()

106 
self.h[0] = last 
107 
self.d[last] = 0 
108 
# Restore invariant by sifting up, then down

109 
pos = self._siftup(0) 
110 
self._siftdown(pos)

111 
# Return smallest element

112 
return elt

113  
114 
def update(self, elt, new): 
115 
"""Replace an element in the queue with a new one."""

116 
# Replace

117 
pos = self.d[elt]

118 
self.h[pos] = new

119 
del self.d[elt] 
120 
self.d[new] = pos

121 
# Restore invariant by sifting up, then down

122 
pos = self._siftup(pos)

123 
self._siftdown(pos)

124  
125 
def remove(self, elt): 
126 
"""Remove an element from the queue."""

127 
# Find and remove element

128 
try:

129 
pos = self.d[elt]

130 
del self.d[elt] 
131 
except KeyError: 
132 
# Not in queue

133 
raise

134 
# If elt is last item, remove and return

135 
if pos == len(self.h)  1: 
136 
self.h.pop()

137 
return

138 
# Replace elt with last element

139 
last = self.h.pop()

140 
self.h[pos] = last

141 
self.d[last] = pos

142 
# Restore invariant by sifting up, then down

143 
pos = self._siftup(pos)

144 
self._siftdown(pos)

145  
146 
def _siftup(self, pos): 
147 
"""Move element at pos down to a leaf by repeatedly moving the smaller

148 
child up."""

149 
h, d = self.h, self.d 
150 
elt = h[pos] 
151 
# Continue until element is in a leaf

152 
end_pos = len(h)

153 
left_pos = (pos << 1) + 1 
154 
while left_pos < end_pos:

155 
# Left child is guaranteed to exist by loop predicate

156 
left = h[left_pos] 
157 
try:

158 
right_pos = left_pos + 1

159 
right = h[right_pos] 
160 
# Outofplace, swap with left unless right is smaller

161 
if right < left:

162 
h[pos], h[right_pos] = right, elt 
163 
pos, right_pos = right_pos, pos 
164 
d[elt], d[right] = pos, right_pos 
165 
else:

166 
h[pos], h[left_pos] = left, elt 
167 
pos, left_pos = left_pos, pos 
168 
d[elt], d[left] = pos, left_pos 
169 
except IndexError: 
170 
# Left leaf is the end of the heap, swap

171 
h[pos], h[left_pos] = left, elt 
172 
pos, left_pos = left_pos, pos 
173 
d[elt], d[left] = pos, left_pos 
174 
# Update left_pos

175 
left_pos = (pos << 1) + 1 
176 
return pos

177  
178 
def _siftdown(self, pos): 
179 
"""Restore invariant by repeatedly replacing outofplace element with

180 
its parent."""

181 
h, d = self.h, self.d 
182 
elt = h[pos] 
183 
# Continue until element is at root

184 
while pos > 0: 
185 
parent_pos = (pos  1) >> 1 
186 
parent = h[parent_pos] 
187 
if parent > elt:

188 
# Swap outofplace element with parent

189 
h[parent_pos], h[pos] = elt, parent 
190 
parent_pos, pos = pos, parent_pos 
191 
d[elt] = pos 
192 
d[parent] = parent_pos 
193 
else:

194 
# Invariant is satisfied

195 
break

196 
return pos
