interactive-coding-challenges/graphs_trees/min_heap/min_heap_solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook was prepared by [Donne Martin](https://github.com/donnemartin). Source and license info is on [GitHub](https://github.com/donnemartin/interactive-coding-challenges)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solution Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem: Implement a min heap with extract_min and insert methods.\n",
    "\n",
    "* [Constraints](#Constraints)\n",
    "* [Test Cases](#Test-Cases)\n",
    "* [Algorithm](#Algorithm)\n",
    "* [Code](#Code)\n",
    "* [Unit Test](#Unit-Test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constraints\n",
    "\n",
    "* Can we assume the inputs are ints?\n",
    "    * Yes\n",
    "* Can we assume this fits memory?\n",
    "    * Yes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Cases\n",
    "\n",
    "* Extract min of an empty tree\n",
    "* Extract min general case\n",
    "* Insert into an empty tree\n",
    "* Insert general case (left and right insert)\n",
    "\n",
    "<pre>\n",
    "          _5_\n",
    "        /     \\\n",
    "       20     15\n",
    "      / \\    /  \\\n",
    "     22  40 25\n",
    "     \n",
    "extract_min(): 5\n",
    "\n",
    "          _15_\n",
    "        /      \\\n",
    "       20      25\n",
    "      / \\     /  \\\n",
    "     22  40 \n",
    "\n",
    "insert(2):\n",
    "\n",
    "          _2_\n",
    "        /     \\\n",
    "       20      5\n",
    "      / \\     / \\\n",
    "     22  40  25  15\n",
    "</pre>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Algorithm\n",
    "\n",
    "A heap is a complete binary tree where each node is smaller than its children.\n",
    "\n",
    "### extract_min\n",
    "\n",
    "<pre>\n",
    "          _5_\n",
    "        /     \\\n",
    "       20     15\n",
    "      / \\    /  \\\n",
    "     22  40 25\n",
    "\n",
    "Save the root as the value to be returned: 5\n",
    "Move the right most element to the root: 25\n",
    "\n",
    "          _25_\n",
    "        /      \\\n",
    "       20      15\n",
    "      / \\     /  \\\n",
    "     22  40 \n",
    "\n",
    "Bubble down 25: Swap 25 and 15 (the smaller child)\n",
    "\n",
    "          _15_\n",
    "        /      \\\n",
    "       20      25\n",
    "      / \\     /  \\\n",
    "     22  40 \n",
    "\n",
    "Return 5\n",
    "</pre>\n",
    "\n",
    "We'll use an array to represent the tree, here are the indices:\n",
    "\n",
    "<pre>\n",
    "          _0_\n",
    "        /     \\\n",
    "       1       2\n",
    "      / \\     / \\\n",
    "     3   4   \n",
    "</pre>\n",
    "\n",
    "To get to a child, we take 2 * index + 1 (left child) or 2 * index + 2 (right child).\n",
    "\n",
    "For example, the right child of index 1 is 2 * 1 + 2 = 4.\n",
    "\n",
    "Complexity:\n",
    "* Time: O(lg(n))\n",
    "* Space: O(lg(n)) for the recursion depth (tree height), or O(1) if using an iterative approach\n",
    "\n",
    "### insert\n",
    "\n",
    "<pre>\n",
    "          _5_\n",
    "        /     \\\n",
    "       20     15\n",
    "      / \\    /  \\\n",
    "     22  40 25\n",
    "\n",
    "insert(2):\n",
    "Insert at the right-most spot to maintain the heap property.\n",
    "\n",
    "          _5_\n",
    "        /     \\\n",
    "       20     15\n",
    "      / \\    /  \\\n",
    "     22  40 25   2\n",
    "\n",
    "Bubble up 2: Swap 2 and 15\n",
    "\n",
    "          _5_\n",
    "        /     \\\n",
    "       20     2\n",
    "      / \\    / \\\n",
    "     22  40 25  15\n",
    "\n",
    "Bubble up 2: Swap 2 and 5\n",
    "\n",
    "          _2_\n",
    "        /     \\\n",
    "       20     5\n",
    "      / \\    / \\\n",
    "     22  40 25  15\n",
    "</pre>\n",
    "\n",
    "We'll use an array to represent the tree, here are the indices:\n",
    "\n",
    "<pre>\n",
    "          _0_\n",
    "        /     \\\n",
    "       1       2\n",
    "      / \\     / \\\n",
    "     3   4   5   6\n",
    "</pre>\n",
    "\n",
    "To get to a parent, we take (index - 1) // 2.  \n",
    "\n",
    "For example, the parent of index 6 is (6 - 1) // 2 = 2.\n",
    "\n",
    "Complexity:\n",
    "* Time: O(lg(n))\n",
    "* Space: O(lg(n)) for the recursion depth (tree height), or O(1) if using an iterative approach"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting min_heap.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile min_heap.py\n",
    "from __future__ import division\n",
    "\n",
    "import sys\n",
    "\n",
    "\n",
    "class MinHeap(object):\n",
    "\n",
    "    def __init__(self):\n",
    "        self.array = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.array)\n",
    "\n",
    "    def extract_min(self):\n",
    "        if not self.array:\n",
    "            return None\n",
    "        if len(self.array) == 1:\n",
    "            return self.array.pop(0)\n",
    "        minimum = self.array[0]\n",
    "        # Move the last element to the root\n",
    "        self.array[0] = self.array.pop(-1)\n",
    "        self._bubble_down(index=0)\n",
    "        return minimum\n",
    "\n",
    "    def peek_min(self):\n",
    "        return self.array[0] if self.array else None\n",
    "\n",
    "    def insert(self, key):\n",
    "        if key is None:\n",
    "            raise TypeError('key cannot be None')\n",
    "        self.array.append(key)\n",
    "        self._bubble_up(index=len(self.array) - 1)\n",
    "\n",
    "    def _bubble_up(self, index):\n",
    "        if index == 0:\n",
    "            return\n",
    "        index_parent = (index - 1) // 2\n",
    "        if self.array[index] < self.array[index_parent]:\n",
    "            # Swap the indices and recurse\n",
    "            self.array[index], self.array[index_parent] = \\\n",
    "                self.array[index_parent], self.array[index]\n",
    "            self._bubble_up(index_parent)\n",
    "\n",
    "    def _bubble_down(self, index):\n",
    "        min_child_index = self._find_smaller_child(index)\n",
    "        if min_child_index == -1:\n",
    "            return\n",
    "        if self.array[index] > self.array[min_child_index]:\n",
    "            # Swap the indices and recurse\n",
    "            self.array[index], self.array[min_child_index] = \\\n",
    "                self.array[min_child_index], self.array[index]\n",
    "            self._bubble_down(min_child_index)\n",
    "\n",
    "    def _find_smaller_child(self, index):\n",
    "        left_child_index = 2 * index + 1\n",
    "        right_child_index = 2 * index + 2\n",
    "        # No right child\n",
    "        if right_child_index >= len(self.array):\n",
    "            # No left child\n",
    "            if left_child_index >= len(self.array):\n",
    "                return -1\n",
    "            # Left child only\n",
    "            else:\n",
    "                return left_child_index\n",
    "        else:\n",
    "            # Compare left and right children\n",
    "            if self.array[left_child_index] < self.array[right_child_index]:\n",
    "                return left_child_index\n",
    "            else:\n",
    "                return right_child_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run min_heap.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unit Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting test_min_heap.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile test_min_heap.py\n",
    "import unittest\n",
    "\n",
    "\n",
    "class TestMinHeap(unittest.TestCase):\n",
    "\n",
    "    def test_min_heap(self):\n",
    "        heap = MinHeap()\n",
    "        self.assertEqual(heap.peek_min(), None)\n",
    "        self.assertEqual(heap.extract_min(), None)\n",
    "        heap.insert(20)\n",
    "        self.assertEqual(heap.array[0], 20)\n",
    "        heap.insert(5)\n",
    "        self.assertEqual(heap.array[0], 5)\n",
    "        self.assertEqual(heap.array[1], 20)\n",
    "        heap.insert(15)\n",
    "        self.assertEqual(heap.array[0], 5)\n",
    "        self.assertEqual(heap.array[1], 20)\n",
    "        self.assertEqual(heap.array[2], 15)\n",
    "        heap.insert(22)\n",
    "        self.assertEqual(heap.array[0], 5)\n",
    "        self.assertEqual(heap.array[1], 20)\n",
    "        self.assertEqual(heap.array[2], 15)\n",
    "        self.assertEqual(heap.array[3], 22)\n",
    "        heap.insert(40)\n",
    "        self.assertEqual(heap.array[0], 5)\n",
    "        self.assertEqual(heap.array[1], 20)\n",
    "        self.assertEqual(heap.array[2], 15)\n",
    "        self.assertEqual(heap.array[3], 22)\n",
    "        self.assertEqual(heap.array[4], 40)\n",
    "        heap.insert(3)\n",
    "        self.assertEqual(heap.array[0], 3)\n",
    "        self.assertEqual(heap.array[1], 20)\n",
    "        self.assertEqual(heap.array[2], 5)\n",
    "        self.assertEqual(heap.array[3], 22)\n",
    "        self.assertEqual(heap.array[4], 40)\n",
    "        self.assertEqual(heap.array[5], 15)\n",
    "        mins = []\n",
    "        while heap:\n",
    "            mins.append(heap.extract_min())\n",
    "        self.assertEqual(mins, [3, 5, 15, 20, 22, 40])\n",
    "        print('Success: test_min_heap')\n",
    "\n",
    "        \n",
    "def main():\n",
    "    test = TestMinHeap()\n",
    "    test.test_min_heap()\n",
    "\n",
    "    \n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success: test_min_heap\n"
     ]
    }
   ],
   "source": [
    "%run -i test_min_heap.py"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}