heap based priority queue
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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