Search code examples
pythonobjectheapmin-heap

min heap in python

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?

Solution

  • 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