min_heap.py

# So, you want to learn MinHeap. Here's an example of min heap.
# Sometimes you panic when you see so many lines of code, but trust me it is worth to read.
# Because there are only three operations: get peek, add item, and poll item.
# The functions are well divided in small functions, very simple, easy to understand.
class MinHeap():
    def __init__(self):
        self.arr = []

    def peek(self):
        if self._is_empty():
            raise ValueError("MinHeap is empty")
        return self.arr[0]

    def poll(self):
        """extract the minimum element from MinHeap"""
        if self._is_empty():
            raise ValueError("MinHeap is Empty")
        if self._size() == 1:
            item = self.arr.pop()
            return item
        item, self.arr[0] = self.arr[0], self.arr.pop()
        self._heapify_down()
        return item

    def add(self, item):
        self.arr.append(item)
        self._heapify_up()

    def _heapify_up(self):
        this = self._size() - 1
        while self._has_parent(this) and self._parent(this) > self.arr[this]:
            self._swap(self._get_parent_index(this), this)
            this = self._get_parent_index(this)

    def _heapify_down(self):
        this = 0
        while self._has_left(this):
            child_index = self._get_left_index(this)
            if self._has_right(this) and \
               self._right(this) < self._left(this):
                child_index = self._get_right_index(this)
            if self.arr[this] < self.arr[child_index]:
                break
            self._swap(this, child_index)
            this = child_index

    def _size(self):
        return len(self.arr)

    def _is_empty(self):
        return self._size() == 0

    def _get_left_index(self, parent_index):
        """ [1, 2, 3]
             |  |
             `parent
                |
                 `left child
        """
        return 2 * parent_index + 1

    def _get_right_index(self, parent_index):
        """ [1, 2, 3]
             |     |
             `parent
                   |
                   `right child
        """
        return 2 * parent_index + 2

    def _get_parent_index(self, child_index):
        return (child_index - 1) // 2

    def _has_left(self, index):
        return self._get_left_index(index) < len(self.arr)

    def _has_right(self, index):
        return self._get_right_index(index) < len(self.arr)

    def _has_parent(self, index):
        return self._get_parent_index(index) < len(self.arr)

    def _left(self, i):
        child_idx = self._get_left_index(i)
        return self.arr[child_idx]

    def _right(self, i):
        child_idx = self._get_right_index(i)
        return self.arr[child_idx]

    def _parent(self, i):
        parent_idx = self._get_parent_index(i)
        return self.arr[parent_idx]

    def _swap(self, i, j):
        self.arr[i], self.arr[j] = self.arr[j], self.arr[i]