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)
--
http://mail.python.org/mailman/listinfo/python-list
