For people who prefer object-oriented style, this wrapper class provides a cleaner interface to the functions in the heapq module. It also provides some additional capabilities, including reduce, which is useful in some graph algorithms.
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
import heapq class Heap(list): """This is a wrapper class for the heap functions provided by the heapq module. """ __slots__ = () def __init__(self, t=): self.extend(t) self.heapify() push = heapq.heappush popmin = heapq.heappop replace = heapq.heapreplace heapify = heapq.heapify def pushpop(self, item): "Push the item onto the heap and then pop the smallest value" if self and self < item: return heapq.heapreplace(self, item) return item def __iter__(self): "Return a destructive iterator over the heap's elements" try: while True: yield self.popmin() except IndexError: pass def reduce(self, pos, newitem): "Replace self[pos] with a lower value item and then reheapify" while pos > 0: parentpos = (pos - 1) >> 1 parent = self[parentpos] if parent <= newitem: break self[pos] = parent pos = parentpos self[pos] = newitem def is_heap(self): "Return True if the heap has the heap property; False otherwise" n = len(self) # The largest index there's any point to looking at # is the largest with a child index in-range, so must have 2*i + 1 < n, # or i < (n-1)/2. If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so # j-1 is the largest, which is n//2 - 1. If n is odd = 2*j+1, this is # (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1. try: for i in xrange(n//2): if self[i] > self[2*i+1]: return False if self[i] > self[2*i+2]: return False except IndexError: pass return True def heapsort(seq): return [x for x in Heap(seq)] if __name__ == '__main__': from random import randint, shuffle # generate a random test case n = 15 data = [randint(1,n) for i in xrange(n)] shuffle(data) print data # test the constructor heap = Heap(data) print heap, heap.is_heap() # test popmin sorted =  while heap: sorted.append(heap.popmin()) data.sort() print heap, heap.is_heap() print data == sorted # test 2 shuffle(data) print data # test push for item in data: heap.push(item) print heap, heap.is_heap() # test __iter__ sorted = [x for x in heap] data.sort() print data == sorted # test 3 shuffle(data) print data heap = Heap(data) print heap, heap.is_heap() # test reduce for i in range(5): pos = randint(0,n-1) decr = randint(1,10) item = heap[pos] - decr heap.reduce(pos, item) # test is_heap heap = Heap(data) count = 0 while 1: shuffle(heap) if heap.is_heap(): print heap break else: count += 1 print 'It took', count, 'tries to find a heap by chance.' print heapsort(data) try: heap.x = 5 except AttributeError: print "Can't add attributes."
A heap, as implemented by the heapq module, is a list that happens to have the heap property, which is that there is no value of k such that heap[k] > heap[2k+1] or heap[k] > heap[2k+2]. The heapq module provides functions to transform a list into a heap and to add and remove elements, among others.
For people who prefer object-oriented style, the Heap class provides a cleaner interface to the functions in the heapq module and provides some additional capabilities.
Heap extends the list class, so Heaps provide all list methods, but methods that modify the list might break the heap property. The is_heap method checks whether a Heap has the heap property.
The methods push, popmin, replace and heapify are just aliases for heappush, heappop, heapreplace and heapify. popmin is so named partly to document its function and partly to avoid conflict with list.pop.
Specifying __slots__ is an optimization that saves memory, but it makes it impossible to add additional attributes to a Heap.
pushpop is a variant of replace that optimizes the case where the item being pushed would immediately be popped.
__iter__ returns a destructive iterator; each time next is invoked, it pops an item from the heap. In this case, the iterator is implemented by a generator.
reduce is sometimes called "reduce-key" in the context of graph algorithms. It replaces the item at the given position with a new item (which must have lower value than the old item) and then restores the heap property.
The test code demonstrates the use of each method.
This class combines ideas of mine with suggestions posted on Python-dev by Raymond Hettinger, Jeremy Fincher and Scott David Daniels. It also contains code I adapted from the Python implementation of the heapq module, which was written by Kevin O'Connor and augmented by Tim Peters and Raymond Hettinger.