min_heap.py

import unittest
from functools import partial


class BinaryTree(object):
    """
    In computer science, a binary tree is a tree data structure in which each 
    node has at most two children (referred to as the left child and the right
    child). 
    
    In a binary tree, the degree of each node can be at most two. 
    
    Binary trees are used to implement binary search trees and binary heaps,
    and are used for efficient searching and sorting. 
    
    A binary tree is a special case of a K-ary tree, where k is 2.
    """
    def __init__(self, data=None):
        self._items = []

        if data is not None:
            self.initialize(data)

    def initialize(self, data):
        for item in data:
            self.insert(item)

    def get_parent_index(self, i):
        return (i+1)//2-1

    @property
    def last_index(self):
        return len(self._items) -1

    def child(self, c):
        if c > self.last_index:
            return None
        return c

    def left_child(self, i):
        return self.child(2*i+1)

    def right_child(self, i):
        return self.child(2*i+2)

    def get_children_index(self, i):
        return [self.left_child(i), self.right_child(i)]

    def swap(self, i, j):
        self._items[i], self._items[j] = self._items[j], self._items[i]


class BinaryHeap(BinaryTree):
    """
    A binary heap is a heap data structure created using a binary tree. 
    It can be seen as a binary tree with two additional constraints:

    * Shape property
    The tree is a complete binary tree; that is, all levels of the tree, 
    except possibly the last one (deepest) are fully filled, and, 
    if the last level of the tree is not complete, 
    the nodes of that level are filled from left to right.

    * Heap property
    All nodes are either [greater than or equal to] or [less than or equal to]
    each of its children, 
    according to a comparison predicate defined for the heap.
    """

    def insert(self, element):
        """
        To add an element to a heap we must perform an up-heap operation 
        by following this algorithm:

        1. Add the element to the bottom level of the heap.
        2. Compare the added element with its parent; if they are in the 
           correct order, stop.
        3. If not, swap the element with its parent and return to the previous
           step.
        """
        self._items.append(element)

        index = self.last_index

        while True:
            parent = self.get_parent_index(index)

            if parent < 0:
                break

            if self._items[parent] <= element:
                break

            self.swap(parent, index)

            index = parent

    def swap_if_needed(self, i, j):
        if j is None:
            return

        if self._items[i] < self._items[j]:
            return

        self.swap(i, j)

        return True

    def flip_the_kids(self, i):
        return map(partial(self.swap_if_needed, i), self.get_children_index(i))

    def delete(self):
        """
        The procedure for deleting the root from the heap (effectively
        extracting the maximum element in a max-heap or the minimum element
        in a min-heap) and restoring the properties is called down-heap.

        1. Replace the root of the heap with the last element on the last level.
        2. Compare the new root with its children; if they are in the correct
           order, stop.
        3. If not, swap the element with one of its children and return to the
           previous step. (Swap with its smaller child in a min-heap and its
           larger child in a max-heap.)
        """
        if not self._items:
            return None

        top = self._items[0]

        self._items[0] = self._items[-1]
        del self._items[-1]

        try:
            for i in xrange(self.last_index+1):
                self.flip_the_kids(i)
        except StopIteration:
            pass

        return top


class BinaryHeapTestCase(unittest.TestCase):
    def assertEntireTree(self, data, expected):
        self.assertEqual(BinaryHeap(data)._items, expected)

    def assertDeleteFromTree(self, data, expected_root, expected):
        b = BinaryHeap(data)
        root = b.delete()

        self.assertEqual(root, expected_root)
        self.assertEqual(b._items, expected)

    def test_init_empty_case(self):
        self.assertEntireTree([], [])
         
    def test_init_one_element(self):
        self.assertEntireTree([1], [1])

    def test_init_two_element(self):
        self.assertEntireTree([1, 1], [1, 1])

    def test_init_two_element_different_prio(self):
        self.assertEntireTree([1, 10], [1, 10])

    def test_init_three_element_different_prio(self):
        self.assertEntireTree([1, 50, 10], [1, 50, 10])

    def test_init_four_element_different_prio(self):
        self.assertEntireTree([1, 50, 10, 100], [1, 50, 10, 100])

    def test_init_lots_element_different_prio(self):
        self.assertEntireTree(
            [1, 50, 10, 100, 2, 5, 3, 15],
            [1, 2, 3, 15, 50, 10, 5, 100]
        )

    def test_delete_from_empty(self):
        self.assertDeleteFromTree([], None, [])

    def test_delete_from_single_item_tree(self):
        self.assertDeleteFromTree([1], 1, [])

    def test_delete_from_two_item_tree(self):
        self.assertDeleteFromTree([1, 2], 1, [2])

    def test_delete_from_three_item_tree(self):
        self.assertDeleteFromTree([1, 2, 3], 1, [2, 3])

    def test_delete_from_four_item_tree(self):
        self.assertDeleteFromTree([2, 3, 4, 5], 2, [3, 5, 4])

    def test_delete_from_five_item_tree(self):
        self.assertDeleteFromTree([2, 3, 5, 4, 8, 7, 6], 2, [3, 4, 5, 6, 8, 7])

    def test_sorting(self):
        data = [10, 14, 4, 13, 12, 1234, 1, 500, 66, 1, 234]
        s = []

        b = BinaryHeap(data)

        while True:
            i = b.delete()
            
            if i is None:
                break

            s.append(i)

        self.assertEqual(s, list(sorted(data)))