priority_queue.py

# 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 HiredInTech.com 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.
        self.heap.append(value)
        # 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
            else:
                break

    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
            else:
                break

        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