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