Heap implement

heap.py

import math
from cStringIO import StringIO


class Heap(object):
    '''                                       index from parent
                        0                     0
                1               2             0*2+1=1  0*2+2=2
            3       4       5       6         1*2+1=3  1*2+2=4 2*2+1=5 2*2+1=6
        7      8 9    10 11   12  13  14      3*2+1=7

        [0,1,2,3,4,5,6,7,8,9,10,11,13,14]
    '''
    min_heap = True

    def __init__(self, min_heap=True):
        self.heap = []
        self.min_heap = min_heap

    @staticmethod
    def parent_index(index):
        if index <= 0:
            return 0
        return int((index - 1) / 2)

    @staticmethod
    def left_child_index(index):
        return index * 2 + 1

    @staticmethod
    def right_child_index(index):
        return index * 2 + 2

    def is_out_of_order(self, child_index, parent_index):
        return (
            self.min_heap and self.heap[parent_index] > self.heap[child_index]
        ) or (
            not self.min_heap and self.heap[parent_index] < self.heap[child_index]
        )

    def bubble_up(self, index):
        ''' Compare the current index with its parent
            swap if wrong order
        '''
        if index <= 0:
            return
        parent_idx = self.parent_index(index)
        if self.is_out_of_order(index, parent_idx):
            self.heap[index], self.heap[parent_idx] = (
                self.heap[parent_idx], self.heap[index])
            self.bubble_up(parent_idx)

    def bubble_down(self, index):
        last_idx = len(self.heap) - 1
        if index >= last_idx:
            return
        left_idx = self.left_child_index(index)
        right_idx = self.right_child_index(index)
        if left_idx > last_idx:
            return
        if (right_idx <= last_idx and self.heap[right_idx] < self.heap[left_idx]
                    and self.is_out_of_order(right_idx, index)):
                self.heap[index], self.heap[right_idx] = (
                    self.heap[right_idx], self.heap[index])
                self.bubble_down(right_idx)
                return
        if self.is_out_of_order(left_idx, index):
            self.heap[index], self.heap[left_idx] = (
                self.heap[left_idx], self.heap[index])
            self.bubble_down(left_idx)


    def insert(self, num):
        self.heap.append(num)
        self.bubble_up(len(self.heap) - 1)

    def pop(self):
        value = None
        if not self.heap:
            return value
        value = self.heap[0]
        if len(self.heap) > 1:
            self.heap[0] = self.heap[-1]
        self.heap.pop()
        self.bubble_down(0)
        return value

    def __iter__(self):
        ''' Return values in order '''
        while self.heap:
            yield self.pop()

    def print_tree(self, total_width=80, fill=' '):
        """ Pretty-print a tree. """
        output = StringIO()
        last_row = -1
        for i, n in enumerate(self.heap):
            if i:
                row = int(math.floor(math.log(i+1, 2)))
            else:
                row = 0
            if row != last_row:
                output.write('\n')
            columns = 2**row
            col_width = int(math.floor((total_width * 1.0) / columns))
            output.write(str(n).center(col_width, fill))
            last_row = row
        print output.getvalue()