Minimum heap implementation in Python

Question

This is my take on implementing a minimum heap using sift-down in Python. Is there any way I can shorten the sift-down code? I tried to not use that many if clauses but I seem to couldn't make it any shorter really.

import math


class MinHeap:
   def __init__(self, arr):
       self.arr = arr
       self.n = len(arr)

    def heapify(self):
        depth = int(math.log2(self.n))
        for i in range((2 ** depth) - 2, -1, -1):
            self.sift_down(i)

    def sift_down(self, i):
        while (2 * i) + 1 < self.n or (2 * i) + 2 < self.n:
            i = self.sift_down_level(i)

    def sift_down_level(self, i):
        left_smaller = right_smaller = False
        if self.is_left_child_exists(i) and self.arr[(2 * i) + 1] < self.arr[i]:
            left_smaller = True
        if self.is_right_child_exists(i) and self.arr[(2 * i) + 2] < self.arr[i]:
            right_smaller = True

        if left_smaller and right_smaller:
            if self.arr[(2 * i) + 1] < self.arr[(2 * i) + 2]:
                self.arr[(2 * i) + 1], self.arr[i] = self.arr[i], self.arr[(2 * i) + 1]
                return (2 * i) + 1
            else:
                self.arr[(2 * i) + 2], self.arr[i] = self.arr[i], self.arr[(2 * i) + 2]
                return (2 * i) + 2
        elif left_smaller:
            self.arr[(2 * i) + 1], self.arr[i] = self.arr[i], self.arr[(2 * i) + 1]
            return (2 * i) + 1
        elif right_smaller:
            self.arr[(2 * i) + 2], self.arr[i] = self.arr[i], self.arr[(2 * i) + 2]
            return (2 * i) + 2
        return -1

    def is_left_child_exists(self, i):
        if (2 * i) + 1 < self.n:
            return True
        return False

    def is_right_child_exists(self, i):
        if (2 * i) + 2 < self.n:
            return True
        return False


nums = [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
heap = MinHeap(nums)
heap.heapify()

Answer 1

You should avoid code repetition: perform the swap at one place only. Avoid recalculating 2 * i + 1 multiple times. And, yes, reduce the number of conditions you check. You only need to check which of the two children (if there are two) has the least value, and then compare that value with the parent's value. That's it.

Some other remarks:

• I would not store the length of the array in a property, because you will then have the overhead of keeping it updated every time a value is pushed or pulled from the heap.
• There is no need to perform a logarithm and exponentiation. Just calculate where the parent is of the last node. The last node is at index len()-1, so its parent is at (len()-1-1)//2 which is len()//2-1.
• The while condition in sift_down should not repeat the check that you are going to do in sift_down_level. Just verify that the return value is not -1.
• And call heapify during initialisation

Without making other changes to your design, we get this:

class MinHeap:
    def __init__(self, arr):
       self.arr = arr
       self.heapify()

    def heapify(self):
        for i in range(len(self.arr) // 2 - 1, -1, -1):
            self.sift_down(i)

    def sift_down(self, i):
        while i != -1:
            i = self.sift_down_level(i)

    def sift_down_level(self, i):
        child = i * 2 + 1
        if child < len(self.arr):
            if child + 1 < len(self.arr) and self.arr[child + 1] < self.arr[child]:
                child += 1
            if self.arr[child] < self.arr[i]:
                self.arr[child], self.arr[i] = self.arr[i], self.arr[child]
                return child
        return -1