diff --git a/graphs_trees/check_balance/check_balance_solution.ipynb b/graphs_trees/check_balance/check_balance_solution.ipynb
index af986fc..d7505f9 100644
--- a/graphs_trees/check_balance/check_balance_solution.ipynb
+++ b/graphs_trees/check_balance/check_balance_solution.ipynb
@@ -60,10 +60,11 @@
     "However, we could check whether the tree is balanced while also checking for the heights.\n",
     "\n",
     "* Base case: If the root is None, return 0\n",
-    "* Check the height of `root.left`, if -1, return -1\n",
-    "* Check the height of `root.right`, if -1, return -1\n",
-    "* If the height differences is greater than 1, return -1\n",
-    "* Otherwise, return 1 + max(left height, right height)\n",
+    "* Recursively check whether the left sub tree is balanced, and get its maximum and minimum height\n",
+    "* Recursively Check whether the right sub tree is balanced, and get its maximum and minimum height\n",
+    "* Calculate the maximum height and minimum height of the current tree\n",
+    "* If both sub-trees are balanced, and the maximum and minimum height of the current tree doesn't differ by more than 1, then the current tree is balanced. Otherwise, it is not\n",
+    "* Return whether the current tree is balanced, and the maximum height and minimum height of the current tree\n",
     "    \n",
     "Complexity:\n",
     "* Time: O(n)\n",
@@ -96,26 +97,19 @@
    },
    "outputs": [],
    "source": [
-    "def check_balance(root):\n",
-    "    if check_height(root) == -1:\n",
-    "        return False\n",
-    "    else:\n",
-    "        return True\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",
+    "    return (balanced, min_h, max_h)\n",
     "\n",
-    "def check_height(root):\n",
-    "    if root is None:\n",
-    "        return 0\n",
-    "    left_height = check_height(root.left)\n",
-    "    if left_height == -1:\n",
-    "        return -1\n",
-    "    right_height = check_height(root.right)\n",
-    "    if right_height == -1:\n",
-    "        return -1\n",
-    "    diff_height = left_height - right_height\n",
-    "    if abs(diff_height) > 1:\n",
-    "        return -1\n",
-    "    else:\n",
-    "        return 1 + max(left_height, right_height)"
+    "def check_balance(root):\n",
+    "    balanced, _, _ = __check_balance__(root)\n",
+    "    return balanced"
    ]
   },
   {
@@ -162,6 +156,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",
diff --git a/graphs_trees/check_balance/test_check_balance.py b/graphs_trees/check_balance/test_check_balance.py
index b63de5c..c53bcf3 100644
--- a/graphs_trees/check_balance/test_check_balance.py
+++ b/graphs_trees/check_balance/test_check_balance.py
@@ -18,6 +18,15 @@ class TestCheckBalance(object):
         insert(node, 10)
         assert_equal(check_balance(node), False)
 
+        node = Node(3)
+        insert(node, 2)
+        insert(node, 1)
+        insert(node, 5)
+        insert(node, 4)
+        insert(node, 6)
+        insert(node, 7)
+        assert_equal(check_balance(node), False)
+
         print('Success: test_check_balance')