heapq.heappush and pop to use key

[jm.suresh@no.spam.gmail.com]

I wanted to have a heap of custom objects, and in different heaps I
wanted to have the weights for my elements differently. So, I modified
the heapq module to accept key arguments also.

The untested code is here. Please let me know if you find any bug or
if there is an easy way to do this.

-
Suresh
---------------------------------------------------------------------------------------------------------------------------------------------
__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace',
'nlargest',
           'nsmallest']

from itertools import islice, repeat, count, imap, izip, tee
from operator import itemgetter
import bisect

pass_it = lambda x: x

def heappush(heap, item, key=pass_it):
    """Push item onto heap, maintaining the heap invariant."""
    heap.append(item)
    _siftdown(heap, 0, len(heap)-1, key=key)

def heappop(heap, key=pass_it):
    """Pop the smallest item off the heap, maintaining the heap
invariant."""
    lastelt = heap.pop()    # raises appropriate IndexError if heap is
empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        _siftup(heap, 0, key)
    else:
        returnitem = lastelt
    return returnitem

def heapreplace(heap, item, key=pass_it):
    """Pop and return the current smallest value, and add the new
item.

    This is more efficient than heappop() followed by heappush(), and
can be
    more appropriate when using a fixed-size heap.  Note that the
value
    returned may be larger than item!  That constrains reasonable uses
of
    this routine unless written as part of a conditional replacement:

        if item > heap[0]:
            item = heapreplace(heap, item)
    """
    returnitem = heap[0]    # raises appropriate IndexError if heap is
empty
    heap[0] = item
    _siftup(heap, 0, key)
    return returnitem

def heapify(x, key=pass_it):
    """Transform list into a heap, in-place, in O(len(heap)) time."""
    n = len(x)
    # Transform bottom-up.  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.
    for i in reversed(xrange(n//2)):
        _siftup(x, i, key)

# 'heap' is a heap at all indices >= startpos, except possibly for
pos.  pos
# is the index of a leaf with a possibly out-of-order value.  Restore
the
# heap invariant.
def _siftdown(heap, startpos, pos, key):
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a
place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if key(parent) <= key(newitem):
            break
        heap[pos] = parent
        pos = parentpos
    heap[pos] = newitem


def _siftup(heap, pos, key):
    endpos = len(heap)
    startpos = pos
    newitem = heap[pos]
    # Bubble up the smaller child until hitting a leaf.
    childpos = 2*pos + 1    # leftmost child position
    while childpos < endpos:
        # Set childpos to index of smaller child.
        rightpos = childpos + 1
        if rightpos < endpos and key(heap[rightpos]) <=
key(heap[childpos]): # **CHANGED**
            childpos = rightpos
        # Move the smaller child up.
        heap[pos] = heap[childpos]
        pos = childpos
        childpos = 2*pos + 1
    # The leaf at pos is empty now.  Put newitem there, and bubble it
up
    # to its final resting place (by sifting its parents down).
    heap[pos] = newitem
    _siftdown(heap, startpos, pos, key)