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). """
        try:
            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?

    Reply
    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.

      Reply
  2. John Miller

    Found your problem:

        def __iter__(self):
            """ Get all elements ordered by asc. priority. """
            return iter(self._heap)
    
    Reply
    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!

      Reply

Leave a Reply

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

 

This site uses Akismet to reduce spam. Learn how your comment data is processed.