Skip to content

Instantly share code, notes, and snippets.

@a3linux
Created June 20, 2018 04:31
Show Gist options
  • Save a3linux/fc0987f8f6192f8f1085a9e90744e37d to your computer and use it in GitHub Desktop.
Save a3linux/fc0987f8f6192f8f1085a9e90744e37d to your computer and use it in GitHub Desktop.
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment