Priority Queue ADT¶
In [1]:
import numpy as np
import matplotlib.pyplot as plt
from heapq import heappush, heappop
push( (priority (float), obj) )
pop() -> obj at the front of the line
In [2]:
pq = []
heappush(pq, (10, "Theo"))
heappush(pq, (20, "Chris"))
heappush(pq, (5, "Celia"))
print(heappop(pq))
print(heappop(pq))
(5, 'Celia') (10, 'Theo')
We desire a data structure that allows us to implement push and pop in $O(\log(N))$, where $N$ is no. of items in priority queue
Heaps¶
In [5]:
class HeapTree(object):
def __init__(self):
self._arr = []
def __len__(self):
return len(self._arr)
def _children(self, i):
"""
Parameters
----------
i: int
Index of a node
Returns
-------
list of indices of children of that node
"""
children = []
if 2*i + 1 < len(self._arr):
children.append(2*i + 1)
if 2*i + 2 < len(self._arr):
children.append(2*i + 2)
return children
def _parent(self, i):
return (i-1)//2
def _swap(self, i, j):
self._arr[i], self._arr[j] = self._arr[j], self._arr[i]
def _heapup(self, i):
"""
Keep bubbling the node at i up until it
satisfies the heap condition
"""
if i > 0:
parent = self._parent(i)
if self._arr[i][0] < self._arr[parent][0]:
self._arr[i], self._arr[parent] = self._arr[parent], self._arr[i]
self._heapup(parent)
def _heapdown(self, i):
children = self._children(i)
if len(children) > 0:
## Make child be the smaller of the two children
child = children[0]
if len(children) > 1:
if self._arr[children[1]][0] < self._arr[children[0]][0]:
child = children[1]
if self._arr[child][0] < self._arr[i][0]:
self._arr[child], self._arr[i] = self._arr[i], self._arr[child]
self._heapdown(child)
def push(self, entry):
"""
Parameters
----------
entry: (priority, obj)
"""
## Step 1: Put this entry at the end of the _arr
self._arr.append(entry)
## Step 2: Bubble up entry until the heap condition is satisfied
self._heapup(len(self._arr)-1)
def pop(self):
assert(len(self._arr) > 0)
ret = self._arr[0][1]
## Move the last element to the root
self._arr[0] = self._arr[-1]
## Take off the last element
self._arr.pop()
## Fix up the internal structure
self._heapdown(0)
return ret
def draw(self, fac=1.5):
N = len(self._arr)
height = int(np.ceil(np.log2(N)))
width = fac*2**height
xs = np.zeros(N)
ys = np.zeros(N)
level = -1
xi = 0
# First draw nodes, and remember positions
# in the process
x0 = width/2
for i in range(N):
if np.log2(i+1) == int(np.log2(i+1)):
level += 1
xi = 0
x0 -= fac*2**(height-level-1)
stride = fac*2**(height-level)
x = x0 + xi*stride
y = -5*level
plt.scatter([x], [y], 100, c='k')
s = "{}".format(self._arr[i][0])
if self._arr[i][1]:
s = s + " ({})".format(self._arr[i][1])
plt.text(x, y-0.7, s, horizontalalignment='center')
xs[i] = x
ys[i] = y
xi += 1
# Next draw edges
for i in range(N):
for j in self._children(i):
plt.plot([xs[i], xs[j]], [ys[i], ys[j]])
plt.axis("off")
plt.axis("equal")
T = HeapTree()
np.random.seed(0)
for x in np.random.randint(0, 100, 13):
T.push((x, x))
T.draw()
## Heap sort! O(N log N) comparison-based sort
## but has bad cache locality
while len(T) > 0:
print(T.pop(), end=",")
9,21,36,44,47,64,67,67,70,83,87,88,88,
In [ ]: