heaps/minheap.py

import math
class minheap(object):
    """
    Heap class - made of keys and items
    methods: build_heap, heappush, heappop
    """

    MIN_HEAP = True

    def __init__(self, nums=None):
        self.heap = []
        if nums:
            self.build_heap(nums)

    def __str__(self):
        return "Min-heap with %s items" % (len(self.heap))

    def max_elements(self):
        return len(self.heap)

    def height(self):
        return math.ceil(math.log(len(self.heap)) / math.log(2))

    def is_leaf(self, i):
        """ returns True if i is a leaf node """
        return i > int(math.ceil((len(self.heap) - 2) / 2.0))

    def parent(self, i):
        if i == 0:
            return -1
        elif i % 2 != 0: # odd
            return (i - 1) / 2
        return (i - 2) / 2

    def leftchild(self, i):
        return 2 * i + 1

    def rightchild(self, i):
        return 2 * i + 2

    def heapify(self, i):
        l = self.leftchild(i)
        r = self.rightchild(i)
        smallest = i
        if l < self.max_elements() and self.heap[l] < self.heap[smallest]:
            smallest = l
        if r < self.max_elements() and self.heap[r] < self.heap[smallest]:
            smallest = r
        if smallest != i:
            self.heap[i], self.heap[smallest] = self.heap[smallest], self.heap[i]
            self.heapify(smallest)

    def build_heap(self, elem):
        """ transforms a list of elements into a heap
        in linear time """
        self.heap = elem[:]
        last_leaf = self.parent(len(self.heap))
        for i in range(last_leaf, -1, -1):
            self.heapify(i)

    def heappush(self, x):
        """ Adds a new item x in the heap"""
        i = len(self.heap)
        self.heap.append(x)
        parent = self.parent(i)
        while parent != -1 and self.heap[i] < self.heap[parent]:
            self.heap[i], self.heap[parent] = self.heap[parent], self.heap[i]
            i = parent
            parent = self.parent(i)

    def heappop(self):
        """ extracts the root of the heap, min or max
        depending on the kind of heap"""
        if self.max_elements():
            self.heap[0], self.heap[-1] = self.heap[-1], self.heap[0]
            pop = self.heap.pop()
            self.heapify(0)
            return pop
        raise Exception("Heap is empty")