Skip to content

Instantly share code, notes, and snippets.

@dcramer
Last active December 11, 2015 04:09
Show Gist options
  • Star 4 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save dcramer/4542897 to your computer and use it in GitHub Desktop.
Save dcramer/4542897 to your computer and use it in GitHub Desktop.
Refined Heap Queue for Python (including capped and max heap implementation)
import heapq
from threading import Lock
class HeapQueue(object):
def __init__(self, values=None, maxsize=None, reversed=False):
"""
Create a new heap queue.
- ``maxsize`` will create a capped queue.
- ``reversed`` will create a max heap queue (default is min).
>>> queue = HeapQueue(maxsize=10, reversed=True)
>>> queue.push('foo', 3)
>>> queue.push('bar', 1)
>>> queue.push('baz', 2)
>>> # default behavior is to exhaust the queue
>>> results = queue.sorted(exhaust=True)
[(3, 'foo'), ('2, 'bar'), (1, 'baz')]
>>> # the queue now has been exhausted
>>> len(queue)
0
"""
self.lock = Lock()
self.lst = values or []
self.maxsize = maxsize
self.reversed = reversed
heapq.heapify(self.lst)
def __len__(self):
return len(self.lst)
def push(self, element, score):
if not self.reversed:
score = score * -1
element = (score, element)
with self.lock:
if self.maxsize and len(self.lst) >= self.maxsize:
heapq.heappushpop(self.lst, element)
else:
heapq.heappush(self.lst, element)
def pop(self):
with self.lock:
score, element = heapq.heappop(self.lst)
if self.reversed:
score = score * -1
return score, element
def sorted(self):
with self.lock:
results = [heapq.heappop(self.lst) for x in xrange(len(self.lst))]
for score, element in reversed(results):
if not self.reversed:
score = score * -1
yield score, element
if __name__ == '__main__':
# test min heap queue
queue = HeapQueue(maxsize=2)
queue.push('foo', 3)
queue.push('bar', 1)
queue.push('baz', 2)
results = list(queue.sorted())
assert len(results) == 2
assert results[0] == (1, 'bar'), results[0]
assert results[1] == (2, 'baz'), results[1]
# test max heap queue
queue = HeapQueue(maxsize=2, reversed=True)
queue.push('foo', 3)
queue.push('bar', 1)
queue.push('baz', 2)
results = list(queue.sorted())
assert len(results) == 2
assert results[0] == (3, 'foo'), results[0]
assert results[1] == (2, 'baz'), results[1]
@RonnyPfannschmidt
Copy link

why the copying iter? thats expensive

@RonnyPfannschmidt
Copy link

max parameter should be called something like reverse, that word is more directly related to the ordering scheme
also use a factor thats either 1 or -1

@dcramer
Copy link
Author

dcramer commented Jan 15, 2013

Copying it so it doesnt exhaust on iter.

I had actually originally called it reverse, but ya probably true.

@RonnyPfannschmidt
Copy link

note that the factor was meant to be used always after setting it in init

@YuppY
Copy link

YuppY commented Jan 16, 2013

Why score argument is required?
Optional key parameter in constructor and score calculation on push would be handier.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment