{ "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: Find the lowest common ancestor of two nodes in a binary tree.\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", "* Is this a binary search tree?\n", " * No\n", "* Can we assume the two nodes are in the tree?\n", " * No\n", "* Can we assume this fits memory?\n", " * Yes" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Test Cases\n", "\n", "
\n",
    "          _10_\n",
    "        /      \\\n",
    "       5        9\n",
    "      / \\      / \\\n",
    "     12  3    18  20\n",
    "        / \\       /\n",
    "       1   8     40\n",
    "
\n", "\n", "* 0, 5 -> None\n", "* 5, 0 -> None\n", "* 1, 8 -> 3\n", "* 12, 8 -> 5\n", "* 12, 40 -> 10\n", "* 9, 20 -> 9\n", "* 3, 5 -> 5" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Algorithm\n", "\n", "* Verify both nodes are in the tree\n", "* Base cases\n", " * If the root is None, return None\n", " * If the root is either node, return the root\n", "* Recursively search left and right\n", "* If the left and right are both nodes\n", " * The return the root\n", "* Else, left or right, whichever is valid\n", "\n", "Complexity:\n", "* Time: O(h), where h is the height of the tree\n", "* Space: O(h), where h is the recursion depth" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Code" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "collapsed": true }, "outputs": [], "source": [ "class Node(object):\n", "\n", " def __init__(self, key, left=None, right=None):\n", " self.key = key\n", " self.left = left\n", " self.right = right\n", "\n", " def __repr__(self):\n", " return str(self.key)" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class BinaryTree(object):\n", "\n", " def lca(self, root, node1, node2):\n", " if None in (root, node1, node2):\n", " return None\n", " if (not self._node_in_tree(root, node1) or\n", " not self._node_in_tree(root, node2)):\n", " return None\n", " return self._lca(root, node1, node2)\n", "\n", " def _node_in_tree(self, root, node):\n", " if root is None:\n", " return False\n", " if root is node:\n", " return True\n", " left = self._node_in_tree(root.left, node)\n", " right = self._node_in_tree(root.right, node)\n", " return left or right\n", "\n", " def _lca(self, root, node1, node2):\n", " if root is None:\n", " return None\n", " if root is node1 or root is node2:\n", " return root\n", " left_node = self._lca(root.left, node1, node2)\n", " right_node = self._lca(root.right, node1, node2)\n", " if left_node is not None and right_node is not None:\n", " return root\n", " else:\n", " return left_node if left_node is not None else right_node" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class LcaResult(object):\n", "\n", " def __init__(self, node, is_ancestor):\n", " self.node = node\n", " self.is_ancestor = is_ancestor\n", "\n", "\n", "class BinaryTreeOptimized(object):\n", "\n", " def lca(self, root, node1, node2):\n", " if root is None:\n", " raise TypeError('root cannot be None')\n", " result = self._lca(root, node1, node2)\n", " if result.is_ancestor:\n", " return result.node\n", " return None\n", "\n", " def _lca(self, curr_node, node1, node2):\n", " if curr_node is None:\n", " return LcaResult(None, is_ancestor=False)\n", " if curr_node is node1 and curr_node is node2:\n", " return LcaResult(curr_node, is_ancestor=True)\n", " left_result = self._lca(curr_node.left, node1, node2)\n", " if left_result.is_ancestor:\n", " return left_result\n", " right_result = self._lca(curr_node.right, node1, node2)\n", " if right_result.is_ancestor:\n", " return right_result\n", " if left_result.node is not None and right_result.node is not None:\n", " return LcaResult(curr_node, is_ancestor=True)\n", " elif curr_node is node1 or curr_node is node2:\n", " is_ancestor = left_result.node is not None or \\\n", " right_result.node is not None\n", " return LcaResult(curr_node, is_ancestor)\n", " else:\n", " return LcaResult(left_result.node if left_result.node is not None \\\n", " else right_result.node, is_ancestor=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Unit Test" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Overwriting test_lca.py\n" ] } ], "source": [ "%%writefile test_lca.py\n", "from nose.tools import assert_equal\n", "\n", "\n", "class TestLowestCommonAncestor(object):\n", "\n", " def test_lca(self):\n", " node10 = Node(10)\n", " node5 = Node(5)\n", " node12 = Node(12)\n", " node3 = Node(3)\n", " node1 = Node(1)\n", " node8 = Node(8)\n", " node9 = Node(9)\n", " node18 = Node(18)\n", " node20 = Node(20)\n", " node40 = Node(40)\n", " node3.left = node1\n", " node3.right = node8\n", " node5.left = node12\n", " node5.right = node3\n", " node20.left = node40\n", " node9.left = node18\n", " node9.right = node20\n", " node10.left = node5\n", " node10.right = node9\n", " root = node10\n", " node0 = Node(0)\n", " binary_tree = BinaryTree()\n", " assert_equal(binary_tree.lca(root, node0, node5), None)\n", " assert_equal(binary_tree.lca(root, node5, node0), None)\n", " assert_equal(binary_tree.lca(root, node1, node8), node3)\n", " assert_equal(binary_tree.lca(root, node12, node8), node5)\n", " assert_equal(binary_tree.lca(root, node12, node40), node10)\n", " assert_equal(binary_tree.lca(root, node9, node20), node9)\n", " assert_equal(binary_tree.lca(root, node3, node5), node5)\n", " print('Success: test_lca')\n", "\n", "\n", "def main():\n", " test = TestLowestCommonAncestor()\n", " test.test_lca()\n", "\n", "\n", "if __name__ == '__main__':\n", " main()" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Success: test_lca\n" ] } ], "source": [ "%run -i test_lca.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.5.0" } }, "nbformat": 4, "nbformat_minor": 0 }