Min-Heap in Python

I recently wanted to implement a small event system where events can have different priorities. So for example the event with highest priority (lowest value) should be handled first.
Python comes with a heapq module which can transform a list into a heap in a way that it stays a list, but fulfills all heap constraints. Nevertheless you might want to wrap the heap like this, so you can do nice stuff:

import heapq
class Heap(object):
    """ A neat min-heap wrapper which allows storing items by priority
        and get the lowest item out first (pop()).
        Also implements the iterator-methods, so can be used in a for
        loop, which will loop through all items in increasing priority order.
        Remember that accessing the items like this will iteratively call
        pop(), and hence empties the heap! """

    def __init__(self):
        """ create a new min-heap. """
        self._heap = []

    def push(self, priority, item):
        """ Push an item with priority into the heap.
            Priority 0 is the highest, which means that such an item will
            be popped first."""
        assert priority >= 0
        heapq.heappush(self._heap, (priority, item))

    def pop(self):
        """ Returns the item with lowest priority. """
        item = heapq.heappop(self._heap)[1] # (prio, item)[1] == item
        return item

    def __len__(self):
        return len(self._heap)

    def __iter__(self):
        """ Get all elements ordered by asc. priority. """
        return self

    def next(self):
        """ Get all elements ordered by their priority (lowest first). """
            return self.pop()
        except IndexError:
            raise StopIteration

With this one could do the following:

h = Heap()
# add some nonsense:
h.push(10, "I'm a large one")
h.push(20, "Actually I'm larger")
h.push(5, "size is not everything")
h.push(0, "which size?")

# get it out in a cool way:
for item in h:
   print item

# prints:
# which size?
# size is not everything
# I'm a large one
# Actually I'm larger

4 thoughts on “Min-Heap in Python

  1. John Miller

    When I try running your code, this is the error I get:

    "TypeError: iter() returned non-iterator of type 'Heap'"

    How can I correct this?

    1. joern Post author

      Which python version do you use and how did you input my code sample? There seems to be a white-space problem with empty lines in WordPress and the code highlighting plugin I use, so maybe the whole problem is that your class definition wasn’t complete? (Should’ve given a syntax error though.)

      I tested the code on py2.6 and 2.7, don’t know about 3.x yet.

  2. John Miller

    Found your problem:

        def __iter__(self):
            """ Get all elements ordered by asc. priority. """
            return iter(self._heap)
    1. joern Post author

      seriously? NO!

      If you do it like that, you just iterate over the list backend of the heap which will not be sorted correctly!


Leave a Reply

Your email address will not be published. Required fields are marked *