Skip to content

Instantly share code, notes, and snippets.

@ptbrowne
Created May 7, 2012 08:36
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save ptbrowne/2626679 to your computer and use it in GitHub Desktop.
heap based priority queue
#!/usr/bin/env python
#-*- coding:utf-8 -*-
#author : Patrick Browne
from time import time
from random import randint
class Heap(list):
"""
Implements a heap-based priority queue.
Heritates from list so I can easily to self[] and len(self)
"""
def __init__(self, iterator=[]):
"""
Iterator is a list of values to insert in the heap at the beginning
"""
list.__init__(self)
self._time = []
for v in iterator:
self.push(v)
@staticmethod
def parent(i):
"""
Helper method to get the parent of a node based on its index in the array
"""
return i - 1 >> 1
def is_greater(self, i, j):
"""
Returns True if node #i has greater priority than node #j
If their priorities are equal, compares the time of entry in the heap
and returns True if #i has been added first
"""
if self[i] == self[j] :
return self._time[i] <= self._time[j]
else:
return self[i] >= self[j]
def swap(self, i, j):
"""
Swaps two nodes in the array and also swaps their date of entry
"""
self[i], self[j] = self[j], self[i]
self._time[i], self._time[j] = self._time[j], self._time[i]
def push(self, node):
"""
Push a node in the heap, keeping the heap property.
Returns the node.
"""
# add the node as a leaf
self.append(node)
self._time.append(int(time()*1000000))
# bubble it up
cur = len(self) - 1
ok = False
while not ok and cur > 0:
# if child > parent
i_parent = self.parent(cur)
if self.is_greater(cur, i_parent):
self.swap(cur, i_parent)
cur = i_parent
else:
ok = True
return self[cur]
def pop(self):
"""
Remove and returns the node with highest priority.
"""
if len(self) == 0:
raise IndexError("index out of range")
# save the root
res = self[0]
# swap with most remote leaf
self.swap(0, len(self) - 1)
# delete it now...
del self[len(self) - 1]
# ...and its time too
del self._time[len(self) - 1]
# now we must bubble down the new root
cur = 0
ok = False
while not ok:
# if greater child > parent, swap
i_left = 2*cur + 1
i_right = 2*cur + 2
if i_right < len(self):
i_greater_child = i_left if self[i_left] > self[i_right] else i_right
elif i_left < len(self):
i_greater_child = i_left
else:
i_greater_child = None
if i_greater_child and self.is_greater(i_greater_child, cur):
self.swap(i_greater_child, cur)
cur = i_greater_child
continue
# if we get to this point, the heap property is now respected
ok = True
return res
if __name__ == "__main__":
heap = Heap()
priorities = (randint(0,100) for _ in range(100000))
for p in priorities:
heap.push(p)
assert all(heap.pop() == p for p in sorted(priorities, reverse=True))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment