{ "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", "
```\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",
"```
" ] }, { "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", "
```\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",
"```
\n", "\n", "We'll use an array to represent the tree, here are the indices:\n", "\n", "
```\n",
"          _0_\n",
"        /     \\\n",
"       1       2\n",
"      / \\     / \\\n",
"     3   4   \n",
"```
\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", "
```\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",
"```
\n", "\n", "We'll use an array to represent the tree, here are the indices:\n", "\n", "
```\n",
"          _0_\n",
"        /     \\\n",
"       1       2\n",
"      / \\     / \\\n",
"     3   4   5   6\n",
"```
\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 }