Priority Queue implementation in Python

Priority Queue is an extension of Queue with following properties

  • each item has a priority associated with it
  • the item with high priority is dequeued before other item with low priority

Binary heap is generally preferred implementation for Priority Queue. It takes time O(log n) to re-heapify the remaining items if the min/max item is removed.The same amount of time O(log n) is taken to insert an new item in the heap.

The key functions to implement in Priority Queue include:

  • insert(val)
  • extractMin()/extractMax()
  • remove(i)
  • getMin()/getMax()
  • changePriority(i, val)

Implementation in Python:

def left_child(k):
    return 2 * k + 1


def right_child(k):
    return 2 * k + 2


def is_leaf(k):
    return 2 * k >= len(mylist)


def parent(k):
    return (k - 1) // 2


def getMin():
    return mylist[0]


def heapify(k):
    if is_leaf(k):
        return

    n = len(mylist)

    # get left and right child index
    left = left_child(k)
    right = right_child(k)

    # find the smallest elements among current node, its left child and right child if exists
    smallest = k
    if left < n and mylist[k] > mylist[left]:
        smallest = left
    if right < n and mylist[smallest] > mylist[right]:
        smallest = right

    # swap current node and the smallest child
    if smallest != k:
        mylist[k], mylist[smallest] = mylist[smallest], mylist[k]
        heapify(smallest)


def shift_up(i):
    while i > 0 and mylist[i] < mylist[parent(i)]:
        mylist[i], mylist[parent(i)] = mylist[parent(i)], mylist[i]
        i = parent(i)


def insert(val):
    mylist.append(val)
    shift_up(len(mylist) - 1)


def extractMin():
    min = mylist[0]
    mylist[0] = mylist[len(mylist) - 1]
    mylist.pop(len(mylist) - 1)
    heapify(0)
    return min


def remove(i):
    mylist[i] = getMin() - 1
    shift_up(i)
    extractMin()


def print_heap():
    for i in range(len(mylist) // 2):
        left = 2 * i + 1
        right = 2 * i + 2
        if left < len(mylist) and right < len(mylist):
            print(
                f"parent: {mylist[i]} left-child: {mylist[left]} right-child: {mylist[right]}"
            )
        elif left < len(mylist):
            print(f"parent: {mylist[i]} left-child: {mylist[left]} right-child: None")
            
mylist = []            

Build the binary heap:

insert(6)
insert(4)
insert(8)
insert(2)
insert(9)
insert(7)
print(mylist)

Output:
[2,4,7,6,9,8]


         4        4         2         2          2
4  -->  /   -->  / \  -->  / \  -->  / \  -->  /    \ 
       6        6   8     4   8     4   8     4     7
                         /         / \       / \   /
                        6         6   9     6   9 8

Insert a new item:

insert(1)
print(mylist)

Output:
[1,4,2,6,9,8,7]

      2             2             1
    /   \         /   \         /   \
   4     7  -->  4     1  -->  4     2
  / \   / \     / \   / \     / \   / \
 6   9 8   1   6   9 8   7   6   9 8   7

Extract the min item:

extractMin()
print(mylist)

Output:
[2,4,7,6,9,8]

      1             7             2
    /   \         /   \         /   \
   4     2  -->  4     2  -->  4     7
  / \   / \     / \   /       / \   / 
 6   9 8   7   6   9  8      6   9 8

Remove the item at specified position:

remove(1)
print(mylist)

Output:
[2,6,7,8,9]

       2                1(2-1)         8              2              2
     /   \            /   \          /   \          /   \          /   \
    4     7  -->     2     7  -->   2     7  -->   8     7  -->   6     7
   / \   /          / \   /        / \            / \            / \
  6   9 8          6   9 8        6   9          6   9          8   9