Skip to content

Instantly share code, notes, and snippets.

@benjamin-hodgson
Last active August 29, 2015 14:00
Show Gist options
  • Save benjamin-hodgson/11172923 to your computer and use it in GitHub Desktop.
Save benjamin-hodgson/11172923 to your computer and use it in GitHub Desktop.
A wrapper of the standard library's heapq, ecapsulated as a class.
import collections.abc
import heapq
class HeapQueue(collections.abc.Container, collections.abc.Sized):
"""
A min-heap class using the standard library's heapq module.
This comes in handy if you need a priority-queue algorithm
but you don't need the heavy machinery of queue.PriorityQueue
(like blocking calls, timeouts, thread-safety, and so on).
"""
def __init__(self, iterable=None):
self._heap = list(iterable) if iterable is not None else []
heapq.heapify(self._heap)
def push(self, item):
"""
Push the value item onto the heap, maintaining the heap invariant.
"""
return heapq.heappush(self._heap, item)
def pop(self):
"""
Pop and return the smallest item from the heap, maintaining the heap invariant.
If the heap is empty, IndexError is raised.
"""
return heapq.heappop(self._heap)
def pushpop(self, item):
"""
Push item on the heap, then pop and return the smallest item from the heap.
The combined action runs more efficiently than a push() followed by a separate call to pop().
"""
return heapq.heappushpop(self._heap, item)
def replace(self, item):
"""
Pop and return the smallest item from the heap, and also push the new item. The heap size doesn’t change.
If the heap is empty, IndexError is raised.
This one step operation is more efficient than a pop() followed by push()
and can be more appropriate when using a fixed-size heap.
The pop/push combination always returns an element from the heap and replaces it with `item`.
The value returned may be larger than the item added.
If that isn’t desired, consider using pushpop() instead.
Its push/pop combination returns the smaller of the two values, leaving the larger value on the heap.
"""
return heapq.heapreplace(self._heap, item)
def peek(self):
"""
Return the smallest item on the heap without removing it.
"""
return self._heap[0]
def __contains__(self, item):
return item in self._heap
def __len__(self):
return len(self._heap)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment