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?
-
For a more confortable snippet - (and Python 3 ready), check my answer at stackoverflow.com/questions/8875706/…jsbueno– jsbueno2013-10-08 14:44:49 +00:00Commented Oct 8, 2013 at 14:44
Add a comment
|
2 Answers
Two options (aside from Devin Jeanpierre's suggestion):
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]The
heapqmodule 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.
3 Comments
tzot
Also, the heapq module is implemented both in C and in Python. Importing heapq uses the C module if available.
John Fouhy
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.
LaC
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.
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
4 Comments
John Fouhy
Note that this won't work in python 3.0, which gets rid of cmp .
Devin Jeanpierre
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.
tzot
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.
Devin Jeanpierre
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.