views:

869

answers:

2
+3  Q: 

min heap in python

Hi -

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?

Thanks,

SetJmp

+5  A: 

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
Devin Jeanpierre
Note that this won't work in python 3.0, which gets rid of __cmp__ .
John Fouhy
No, it won't work in Python 3.0, but that doesn't matter, because Python 3.0 lacks the entire concept of "comparator function". The question simply doesn't apply in that case.
Devin Jeanpierre
Um. You should define __le__ rather than __cmp__. That would keep the Python 2/3 compatibility. After all, heapq operations and sort operations use <= as their comparison function.
ΤΖΩΤΖΙΟΥ
I don't care about python 3.0 compatibility. Python 3.0 is still not commonly used, and writing 2/3 quines is difficult and ultimately futile. The most pythonic way to define this in 2.x is as I wrote above. Python 3.0 is backwards incompatible, and what's best in 2.x may not work. Too bad.
Devin Jeanpierre
+7  A: 

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.

John Fouhy
+1: Decorate your data.
S.Lott
Also, the heapq module is implemented both in C and in Python. Importing heapq uses the C module if available.
ΤΖΩΤΖΙΟΥ
Ah, you're right. Well, the important thing for my second suggestion is that python code is available in your lib directory if you want to roll your own.
John Fouhy
I think this is the better solution. It actually takes less code, and is probably faster, since you're using the built-in tuple type instead of defining your own class.
LaC