heap.py

#! /usr/bin/env python

class MinHeap(object):
    """A min heap with a zero based index
    """
    def __init__(self, arr=[]):
        self.arr = arr
        self.size = 0
        self.build_min_heap()

    def min_heapify(self, i):
        """
        Running time: O(lg n)
        """
        l = self._left(i)
        r = self._right(i)
        if l <= self.size-1 and self.arr[l] < self.arr[i]:
            smallest = l
        else:
            smallest = i
        if r <= self.size-1 and self.arr[r] < self.arr[smallest]:
            smallest = r
        if smallest != i:
            self._swap(i, smallest)
            self.min_heapify(smallest)

    def build_min_heap(self):
        """Start with the parents of the leaves
        needed for heapsort

        Running time: O(n)
        """
        self.size = len(self.arr)
        for i in range((self.size-1)//2, -1, -1):
            self.min_heapify(i)

    def check_ri(self):
        """Check representation invariant:
        the min heap property
        """
        bad_nodes = []
        for i in range(self.size):
            p = self._parent(i)
            if not self.arr[p] <= self.arr[i]:
                bad_nodes.append(i)
        if len(bad_nodes) > 0:
            return {False: bad_nodes}
        else:
            return {True}

    def extract_min(self):
        """Removes and returns the element of the arr with the smallest key.

        Running time: O(lg n)
        """
        if self.size < 0:
            raise ValueError('Heap underflow')
        min = self.arr[0]
        self.arr[0] = self.arr[self.size-1]
        self.size -= 1
        self.min_heapify(0)
        return min

    def insert(self, key):
        """
        For sentinal the value infinity either use
            None which is less than any integer or
            float('inf') a special value in python

        Running time: O(lg n)
        """
        self.size += 1
        self.arr.append(float('inf'))
        self._decrease_key(self.size-1, key)

    def min(self):
        return self.arr[0]

    def sort(self):
        """Sorts array in descending order, in place.

        Running time: O(n lg n)
        """
        #sorted_arr = []
        self.build_min_heap()
        # from i=n to 1
        for i in range(self.size-1, 0, -1):
            self.swap(i, 0)
            self.size -= 1
            self.min_heapify(0)
        return self.arr

    def _decrease_key(self, i, key):
        """Decreases value of element i's key to the new value k,
        which is assumed to be at most as big as i's current key value.

        Running time: O(lg n)
        """
        if key > self.arr[i]:
            raise ValueError('new key is larger than current key')
        self.arr[i] = key
        while i > 0 and self.arr[self._parent(i)] > self.arr[i]:
            self._swap(i, self._parent(i))
            i = self._parent(i)

    def _parent(self, i):
        #return (i-0)//2
        return (i-1)//2

    def _left(self, i):
        #return 2*i + 0
        return 2*i + 1

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

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

heap_test.py

#!/usr/bin/env python

import unittest
import sys
import glob
import re
from heap import *

class MinHeapTest(unittest.TestCase):
    def setUp(self):
        self.heap = MinHeap()
        min_heap = [1, 1, 3, 1, 2, 3]
        self.heap.arr = min_heap
        self.heap.size = len(min_heap)

    def test_extract_min(self):
        self.heap.insert(0)
        self.assertEqual(self.heap.extract_min(), 0)

    def test_insert_empty(self):
        self.heap = MinHeap()
        self.heap.insert(8)
        self.assertEqual(self.heap.arr, [8])

    def test_insert_small(self):
        self.heap.insert(0)
        self.assertEqual(self.heap.arr, [0, 1, 1, 1, 2, 3, 3])

    def test_insert_mid(self):
        self.heap.insert(2)
        self.assertEqual(self.heap.arr, [1, 1, 2, 1, 2, 3, 3])

    def test_insert_large(self):
        self.heap.insert(8)
        self.assertEqual(self.heap.arr, [1, 1, 3, 1, 2, 3, 8])

if __name__ == '__main__':
    unittest.main()