Welcome, guest | Sign In | My Account | Store | Cart

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.

Python, 131 lines
  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[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.

2 comments

Peter Fein 16 years ago  # | flag

borrowed methods should be lambdas. Looks like you want this:

push = lambda self, x: heapq.heappush(self, x)

popmin = lambda self: heapq.heappop(self)

replace = lambda self, x: heapq.heapreplace(self, x)

heapify = lambda self: heapq.heapify(self)
Allen Downey (author) 11 years, 6 months ago  # | flag

Peter, your suggestion sounds reasonable, but I don't know why the version using lambdas is better. Can you explain?