Entry

Is there a heap anywhere in the python modules?

May 14th, 2000 05:48
unknown unknown, François Pinard

Here is a draft for a simple Heap module.  This one has been tested.  
See the `test()' function near the end to see an example of usage.


"""\
Handle priority heaps.

Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for all 
k, counting elements from 0.  For the sake of comparison, unexisting 
elements are considered to be infinite. The interesting property of a 
heap is that a[0] is always its smallest element.
"""

class Heap:

    def __init__(self, compare=cmp):
        """\
Set a new heap.  If COMPARE is given, use it instead of built-in 
comparison.

COMPARE, given two items, should return negative, zero or positive 
depending on the fact the first item compares smaller, equal or greater 
than the second item.
"""
        self.compare = compare
        self.array = []

    def __call__(self):
        """\
A heap instance, when called as a function, return its smallest item.
"""
        return self.array[0]

    def __len__(self):
        """\
Return the number of items in the current heap instance.
"""
        return len(self.array)

    def push(self, item):
        """\
Add ITEM to the current heap instance.
"""
        array = self.array
        compare = self.compare
        array.append(item)
        high = len(array) - 1
        while high > 0:
            low = (high-1)/2
            if compare(array[low], array[high]) <= 0:
                break
            array[low], array[high] = array[high], array[low]
            high = low

    def pop(self):
        """\
Remove and return the smallest item from the current heap instance.
"""
        array = self.array
        compare = self.compare
        item = array[0]
        if len(array) == 1:
            del array[0]
        else:
            array[0] = array.pop()
            low, high = 0, 1
            while high < len(array):
                if ((high+1 < len(array)
                     and compare(array[high], array[high+1]) > 0)):
                    high = high+1
                if compare(array[low], array[high]) <= 0:
                    break
                array[low], array[high] = array[high], array[low]
                low, high = high, 2*high+1
        return item

def test(n=2000):
    heap = Heap()
    for k in range(n-1, -1, -1):
        heap.push(k)
    for k in range(n):
        assert k+len(heap) == n
        assert k == heap.pop()