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.
Download
Copy to clipboard1 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[0] < 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.
borrowed methods should be lambdas. Looks like you want this:
Peter, your suggestion sounds reasonable, but I don't know why the version using lambdas is better. Can you explain?