11

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?

Cristian Ciupitu
  • 20,270
  • 7
  • 50
  • 76
Setjmp
  • 27,279
  • 27
  • 74
  • 92
  • For a more confortable snippet - (and Python 3 ready), check my answer at http://stackoverflow.com/questions/8875706/python-heapq-with-custom-compare-predicate/8875823#8875823 – jsbueno Oct 08 '13 at 14:44

2 Answers2

16

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
  • 41,203
  • 19
  • 62
  • 77
  • 1
    Also, the heapq module is implemented both in C and in Python. Importing heapq uses the C module if available. – tzot Mar 25 '09 at 08:59
  • 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 Mar 25 '09 at 21:45
  • 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 Feb 25 '10 at 00:42
14

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
  • 92,913
  • 4
  • 55
  • 79
  • 2
    Note that this won't work in python 3.0, which gets rid of __cmp__ . – John Fouhy Mar 25 '09 at 02:00
  • 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 Mar 25 '09 at 02:12
  • 6
    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. – tzot Mar 25 '09 at 08:58
  • 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 Mar 25 '09 at 11:27