Fix check_balance solution for a new test case

This commit is contained in:
Xiaojian Wang 2015-10-09 22:24:01 -07:00
parent 6276624914
commit edfbec618f
2 changed files with 35 additions and 23 deletions

View File

@ -60,10 +60,11 @@
"However, we could check whether the tree is balanced while also checking for the heights.\n", "However, we could check whether the tree is balanced while also checking for the heights.\n",
"\n", "\n",
"* Base case: If the root is None, return 0\n", "* Base case: If the root is None, return 0\n",
"* Check the height of `root.left`, if -1, return -1\n", "* Recursively check whether the left sub tree is balanced, and get its maximum and minimum height\n",
"* Check the height of `root.right`, if -1, return -1\n", "* Recursively Check whether the right sub tree is balanced, and get its maximum and minimum height\n",
"* If the height differences is greater than 1, return -1\n", "* Calculate the maximum height and minimum height of the current tree\n",
"* Otherwise, return 1 + max(left height, right height)\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", " \n",
"Complexity:\n", "Complexity:\n",
"* Time: O(n)\n", "* Time: O(n)\n",
@ -96,26 +97,19 @@
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"def check_balance(root):\n", "def __check_balance__(root):\n",
" if check_height(root) == -1:\n", " if not root:\n",
" return False\n", " return (True, 0, 0)\n",
" else:\n", " left_balanced, left_min_h, left_max_h = __check_balance__(root.left)\n",
" return True\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", "\n",
"def check_height(root):\n", "def check_balance(root):\n",
" if root is None:\n", " balanced, _, _ = __check_balance__(root)\n",
" return 0\n", " return balanced"
" 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)"
] ]
}, },
{ {
@ -162,6 +156,15 @@
" insert(node, 10)\n", " insert(node, 10)\n",
" assert_equal(check_balance(node), False)\n", " assert_equal(check_balance(node), False)\n",
"\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", " print('Success: test_check_balance')\n",
"\n", "\n",
"\n", "\n",

View File

@ -18,6 +18,15 @@ class TestCheckBalance(object):
insert(node, 10) insert(node, 10)
assert_equal(check_balance(node), False) 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') print('Success: test_check_balance')