Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
# Implementation of a priority queue using a binary heap where the
# binary heap is stored using a vector. The heap stores the minimal
# element at the top, which makes it a min-heap.
# This is sample class for the tutorials.
class PriorityQueue:
def __init__(self):
# We will use a `list` to store the heap values
self.heap = []
def get_size(self):
return len(self.heap)
def insert(self, value):
# Add the new value at the end of the tree, which is
# the first free position in the list.
# Start bubbling up the node if it has a lower value
# then its current parent node.
idx = len(self.heap) - 1
# While the node has not reached the root of the tree
while idx > 0:
# Get the index of the parent node
parent_idx = (idx - 1) // 2
# Check if the current parent has a lower value.
# If so, swap the nodes. Otherwise, stop here.
if self.heap[idx] < self.heap[parent_idx]:
self._swap_nodes(idx, parent_idx)
idx = parent_idx
def get_minimum(self):
# The element at position 0 contains the root node's value.
return self.heap[0]
def pop(self):
# Remove and store the root value to return it in the end
result = self.heap[0]
# If the heap has only one element just pop it and return
if len(self.heap) == 1:
return self.heap.pop()
# Move the last node to the root
self.heap[0] = self.heap.pop()
# Start bubbling down the node we put at the root
n = len(self.heap)
idx = 0
# While the current node has at least one child node,
# check if they must be swapped
while 2 * idx + 1 < n:
child_idx = 2 * idx + 1
child_val = self.heap[child_idx]
# If there is a right child check if it has a smaller value
# than its left sibling.
if child_idx + 1 < n and child_val > self.heap[child_idx + 1]:
child_val = self.heap[child_idx + 1]
child_idx += 1
# Check if the current node has a child with a smaller value.
# If so, swap them. Otherwise, stop iterating down the tree.
if child_val < self.heap[idx]:
self._swap_nodes(idx, child_idx)
idx = child_idx
return result
# A helper method to swap two nodes in the binary heap.
def _swap_nodes(self, idx1, idx2):
tmp = self.heap[idx1]
self.heap[idx1] = self.heap[idx2]
self.heap[idx2] = tmp
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment