min heap in python

Question

I'd like to store a set of objects in a min heap by defining a custom comparison function. I see there is a heapq module available as part of the python distribution. Is there a way to use a custom comparator with this module? If not, has someone else built a custom min heap?

Answer 1

Two options (aside from Devin Jeanpierre's suggestion):

1. Decorate your data before using the heap. This is the equivalent of the key= option to sorting. e.g. if you (for some reason) wanted to heapify a list of numbers according to their sine:

   data = [ # list of numbers ]
   heap = [(math.sin(x), x) for x in data]
   heapq.heapify(heap)
   # get the min element
   item = heappop(heap)[1]
   

2. The heapq module is implemented in pure python. You could just copy it to your working directory and change the relevant bits. From a quick look, you would have to modify siftdown() and siftup(), and possibly nlargest and nsmallest if you need them.

Answer 2

Yes, there is a way. Define a wrapping class that implements your custom comparator, and use a list of those instead of a list of your actual objects. That's about the best there is while still using the heapq module, since it provides no key= or cmp= arguments like the sorting functions/methods do.

def gen_wrapper(cmp):
    class Wrapper(object):
        def __init__(self, value): self.value = value
        def __cmp__(self, obj): return cmp(self.value, obj.value)
    return Wrapper