diff --git a/graphs_trees/check_balance/check_balance_challenge.ipynb b/graphs_trees/check_balance/check_balance_challenge.ipynb
index 3f13284..bcd1fb4 100644
--- a/graphs_trees/check_balance/check_balance_challenge.ipynb
+++ b/graphs_trees/check_balance/check_balance_challenge.ipynb
@@ -36,7 +36,11 @@
     "\n",
     "* Is a balanced tree one where the heights of two sub trees of any node doesn't differ by more than 1?\n",
     "    * Yes\n",
+    "* If this is called on a None input, should we return False?\n",
+    "    * Yes\n",
     "* Can we assume we already have a Node class with an insert method?\n",
+    "    * Yes\n",
+    "* Can we assume this fits memory?\n",
     "    * Yes"
    ]
   },
@@ -46,6 +50,8 @@
    "source": [
     "## Test Cases\n",
     "\n",
+    "* None -> No\n",
+    "* 1 -> Yes\n",
     "* 5, 3, 8, 1, 4 -> Yes\n",
     "* 5, 3, 8, 9, 10 -> No"
    ]
@@ -120,7 +126,11 @@
     "class TestCheckBalance(object):\n",
     "\n",
     "    def test_check_balance(self):\n",
+    "        assert_equal(check_balance(None), False)\n",
+    "\n",
     "        node = Node(5)\n",
+    "        assert_equal(check_balance(node), True)\n",
+    "\n",
     "        insert(node, 3)\n",
     "        insert(node, 8)\n",
     "        insert(node, 1)\n",
@@ -134,6 +144,15 @@
     "        insert(node, 10)\n",
     "        assert_equal(check_balance(node), False)\n",
     "\n",
+    "        node = Node(3)\n",
+    "        insert(node, 2)\n",
+    "        insert(node, 1)\n",
+    "        insert(node, 5)\n",
+    "        insert(node, 4)\n",
+    "        insert(node, 6)\n",
+    "        insert(node, 7)\n",
+    "        assert_equal(check_balance(node), False)\n",
+    "\n",
     "        print('Success: test_check_balance')\n",
     "\n",
     "\n",
@@ -158,21 +177,21 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.10"
+   "pygments_lexer": "ipython3",
+   "version": "3.5.0"
   }
  },
  "nbformat": 4,
diff --git a/graphs_trees/check_balance/check_balance_solution.ipynb b/graphs_trees/check_balance/check_balance_solution.ipynb
index d7505f9..775b786 100644
--- a/graphs_trees/check_balance/check_balance_solution.ipynb
+++ b/graphs_trees/check_balance/check_balance_solution.ipynb
@@ -35,7 +35,11 @@
     "\n",
     "* Is a balanced tree one where the heights of two sub trees of any node doesn't differ by more than 1?\n",
     "    * Yes\n",
+    "* If this is called on a None input, should we return False?\n",
+    "    * Yes\n",
     "* Can we assume we already have a Node class with an insert method?\n",
+    "    * Yes\n",
+    "* Can we assume this fits memory?\n",
     "    * Yes"
    ]
   },
@@ -45,6 +49,8 @@
    "source": [
     "## Test Cases\n",
     "\n",
+    "* None -> No\n",
+    "* 1 -> Yes\n",
     "* 5, 3, 8, 1, 4 -> Yes\n",
     "* 5, 3, 8, 9, 10 -> No"
    ]
@@ -97,18 +103,20 @@
    },
    "outputs": [],
    "source": [
-    "def __check_balance__(root):\n",
+    "def _check_balance(root):\n",
     "    if not root:\n",
     "        return (True, 0, 0)\n",
-    "    left_balanced, left_min_h, left_max_h = __check_balance__(root.left)\n",
-    "    right_balanced, right_min_h, right_max_h = __check_balance__(root.right)\n",
-    "    min_h = min(left_min_h, right_min_h) + 1\n",
-    "    max_h = max(left_max_h, right_max_h) + 1\n",
-    "    balanced = left_balanced and right_balanced and abs(max_h - min_h) <= 1\n",
+    "    left_balanced, left_min_h, left_max_h = _check_balance(root.left)\n",
+    "    right_balanced, right_min_h, right_max_h = _check_balance(root.right)\n",
+    "    min_h = 1 + min(left_min_h, right_min_h)\n",
+    "    max_h = 1 + max(left_max_h, right_max_h)\n",
+    "    balanced = left_balanced and right_balanced and abs(max_h-min_h) <= 1\n",
     "    return (balanced, min_h, max_h)\n",
     "\n",
     "def check_balance(root):\n",
-    "    balanced, _, _ = __check_balance__(root)\n",
+    "    if root is None:\n",
+    "        return False\n",
+    "    balanced, _, _ = _check_balance(root)\n",
     "    return balanced"
    ]
   },
@@ -142,7 +150,11 @@
     "class TestCheckBalance(object):\n",
     "\n",
     "    def test_check_balance(self):\n",
+    "        assert_equal(check_balance(None), False)\n",
+    "\n",
     "        node = Node(5)\n",
+    "        assert_equal(check_balance(node), True)\n",
+    "\n",
     "        insert(node, 3)\n",
     "        insert(node, 8)\n",
     "        insert(node, 1)\n",
@@ -199,21 +211,21 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.10"
+   "pygments_lexer": "ipython3",
+   "version": "3.5.0"
   }
  },
  "nbformat": 4,
diff --git a/graphs_trees/check_balance/test_check_balance.py b/graphs_trees/check_balance/test_check_balance.py
index c53bcf3..843ba6c 100644
--- a/graphs_trees/check_balance/test_check_balance.py
+++ b/graphs_trees/check_balance/test_check_balance.py
@@ -4,7 +4,11 @@ from nose.tools import assert_equal
 class TestCheckBalance(object):
 
     def test_check_balance(self):
+        assert_equal(check_balance(None), False)
+
         node = Node(5)
+        assert_equal(check_balance(node), True)
+
         insert(node, 3)
         insert(node, 8)
         insert(node, 1)