Binary Heap implementation in python

As a follow up to my last article here is the sample implementation of a max heap, that means you will always have a maximum value at the root node.
With simple changes one can reverse this behavior if required.

All code and unit tests are available on github. Here
So without further ado lets dive in.

Our Binary heap class has a constructor that initializes the heap.
The class exposes the following methods

push() 
    pop()
    get_size()
    print_heap() #just for visualization

In the constructor we push the nodes into the heap while calling percolate up function recursively until the heap is balanced, this ensures that the heap is balanced as explained in the previous article

The same process happens when we push new nodes onto the heap.

Here is a snippet of percolate_up method.

def __percolate_up(self, node_to_percolate=0):
        if node_to_percolate == 0:
            current_node = self.heap.index(self.heap[-1])  # node being percolated
        else:
            current_node = node_to_percolate
        print(f"current node : {current_node}")
        # if its the first element we don't need to percolate
        if len(self.heap) == 1:
            print("nothing to percolate up")
            return
        print("percolating..")
        # Using integer division would be the same as (x -1)/2 which allows us to find the parent node
        parent_index = (current_node // 2)
        parent_value = self.heap[parent_index]
        if self.heap[current_node] > parent_value:
            print("swapping")
            # do the swap
            tmp = self.heap[parent_index]
            self.heap[parent_index] = self.heap[current_node]
            self.heap[current_node] = tmp
            # current node takes the index of the parent
            current_node = parent_index
            # swap complete
            print(f"swapped {current_node} to position of {self.heap[parent_index]}")
            return self.__percolate_up(current_node)

Removing nodes

This is done by removing the first node from the array and swapping the last node to the first position, then calling percolate_down function until the heap is balanced.

The implementation of percolate_down is shown below.

def __percolate_down(self, node_to_percolate=0):
        print("Percolating down")
        current_node = node_to_percolate
        print(f"current node {current_node}")
        if not current_node:
            last_node = self.heap.pop()
            # Move last node to be the first
            self.heap = [last_node] + self.heap
        max_child, child_index = self.__max_child(current_node)
        print(f"child {max_child} , child index: {child_index}")
        if not max_child or not child_index:
            return
        # check if swap is necessary
        if max_child >= self.heap[current_node]:
            # do the swap
            tmp = self.heap[current_node]
            self.heap[current_node] = self.heap[child_index]
            self.heap[child_index] = tmp
            # swap complete
            print(f"swapped {self.heap[child_index]} to position of {self.heap[current_node]}")
            # continue percolating at the child level
            return self.__percolate_down(child_index)
    def __max_child(self, node_index):
        left_child_index = (2 * node_index) + 1
        right_child_index = (2 * node_index) + 2
        try:
            left_child = self.heap[left_child_index]
            right_child = self.heap[right_child_index]
            max_child = max(left_child, right_child)
            if left_child == max_child:
                return max_child, left_child_index
            if right_child == max_child:
                return max_child, right_child_index

        except IndexError:
            # means no children for a node exist
            return None, None

Well that’s it, now we know how to implement a maxHeap.

I am sure there are better implementations out there, feel free to suggest improvements via a PR.

Have a beautiful day.