The Heap class is an improvement over a previous recipe (http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/437116) that also wraps the heapq module. There are two main differences from the other recipe: - All methods on this heap preserve the heap property invariant; therefore there is no need for is_heap(). - When creating a new heap, an optional 'key' argument can be specified to determine the comparison key for the items to be pushed into the heap.
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 | import heapq
import itertools as it
from listmixin import ListMixin # (recipe 440656)
__all__ = ['Heap']
class Heap(ListMixin):
'''A list that maintains the heap invariant.'''
def __init__(self, iterable=(), key=None):
'''
@param iterable: An iterable over items to be added to the heap.
@param key: Specifies a function of one argument that is used to
extract a comparison key from each heap element.
'''
self._key = key
self._lst = []
self.extend(iterable)
heapq.heapify(self._lst)
def push(self, item):
'''Push the item onto the heap.'''
return heapq.heappush(self._lst, self._wrap(item))
def popmin(self):
'''Pop the smallest item off the heap'''
return self._unwrap(heapq.heappop(self._lst))
def replace(self, item):
'''Equivalent to "x = heap.popmin(); heap.push(); return x" but more
efficient.
'''
return self._unwrap(heapq.heapreplace(self._lst, self._wrap(item)))
def pushpop(self, item):
'Equivalent to "heap.push(); return heap.popmin()" but more efficient.'
if self and self[0] < item:
return self.replace(item)
return item
def iterpop(self):
'''Return a destructive iterator over the heap's elements.
Each time next is invoked, it pops the smallest item from the heap.
'''
while self:
yield self.popmin()
#---- overrided ListMixin methods ----------------------------------------
def constructor(self, iterable):
return self.__class__(iterable, self._key)
def __len__(self):
return len(self._lst)
def get_element(self, pos):
return self._unwrap(self._lst[pos])
def __setitem__(self, pos, item):
if isinstance(pos, slice):
raise TypeError('Heap objects do no support slice setting')
pos = self._fix_index(pos)
item = self._wrap(item)
lst = self._lst
current = lst[pos]
lst[pos] = item
if item > current: # re-establish the heap invariant
heapq._siftup(lst, pos)
if lst[pos] != item: # item found its way below pos
return
while pos > 0:
parentpos = (pos - 1) >> 1
parent = lst[parentpos]
if parent <= item:
break
lst[pos] = parent
pos = parentpos
lst[pos] = item
def __delitem__(self, pos):
if isinstance(pos, slice):
raise TypeError('Heap objects do no support slice deleting')
pos = self._fix_index(pos)
lst = self._lst
if pos == len(lst)-1:
del lst[-1]
else:
self[pos] = self.pop()
def __iter__(self):
return it.imap(self._unwrap, self._lst)
def __nonzero__(self):
return bool(self._lst)
def __cmp__(self, other):
raise TypeError('Heap objects do not support comparison')
def __eq__(self, other):
if not isinstance(other,Heap) or len(self) != len(other):
return False
for i,j in it.izip(self,other):
if i != j: return False
return True
def __ne__(self,other):
return not self==other
def count(self, item):
return self._lst.count()
append = push
def insert(self, pos, item):
# ignore the position since it's not certain that it can preserve the
# heap invariant
self.push(item)
def extend(self, other):
if self._key is not None:
other = it.imap(self._wrap, other)
push = heapq.heappush; lst = self._lst
for item in other:
push(lst,item)
def sort(self):
lst = self._lst; pop = heapq.heappop
sorted = []; append = sorted.append
while lst:
append(pop(lst))
self._lst = sorted
def reverse(self):
raise TypeError('Heap objects do not support reversal')
#---- 'private' methods --------------------------------------------------
def _wrap(self, item):
if self._key is not None:
item = (self._key(item),item)
return item
def _unwrap(self, item):
if self._key is not None:
item = item[1]
return item
|
Interface: Heap instances support most list methods except for those that are error prone or don't make sense for heaps. For example, setting a single element is supported but setting a slice is not. The reason is that setting a single element, e.g. heap[3] = 'a', means "replace heap[3] with 'a' and restore the heap invariant". As a result, more than one cells are modified in general, not just the one at the specified index. A slice assignment, e.g. heap[3:5] = 'a','b', should be equivalent to "heap[3]='a'; heap[4]='b'". However, heap[4] has most likely changed after the first assignment and it won't probably do what the client meant.
Implementation: The implementation is based on the list mixin recipe (http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/440656). However it overrides __setitem__() and __detitem__() completely instead of implementing set_element() and resize_region(), as the mixin normally requires, for the reason mentioned above.
Testing: Passes successfully a unit test of 34 test cases (not included).