HEAP/priority queue

[François Pinard]

Courageous <jkraska1 at san.rr.com> writes:

> Is there a heap anywhere in the python modules?

Here is a second 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()

-- 
François Pinard   http://www.iro.umontreal.ca/~pinard