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

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.

Python, 148 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
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).