019f8fd211
git-subtree-dir: users/wpcarro git-subtree-mainline:464bbcb15cgit-subtree-split:24f5a642afChange-Id: I6105b3762b79126b3488359c95978cadb3efa789
126 lines
3.7 KiB
Python
126 lines
3.7 KiB
Python
import unittest
|
|
from collections import deque
|
|
|
|
|
|
# is_balanced :: Node(a) -> Bool
|
|
def is_balanced(node):
|
|
q = deque()
|
|
q.append((0, node))
|
|
mn, mx = None, None
|
|
|
|
while q:
|
|
depth, node = q.popleft()
|
|
# Current node is a leaf node
|
|
if not node.left and not node.right:
|
|
mx = depth if mx is None else max(mx, depth)
|
|
mn = depth if mn is None else min(mn, depth)
|
|
if mx - mn > 1:
|
|
return False
|
|
if node.left:
|
|
q.append((depth + 1, node.left))
|
|
if node.right:
|
|
q.append((depth + 1, node.right))
|
|
|
|
return mx - mn <= 1
|
|
|
|
|
|
# Tests
|
|
class Test(unittest.TestCase):
|
|
class BinaryTreeNode(object):
|
|
def __init__(self, value):
|
|
self.value = value
|
|
self.left = None
|
|
self.right = None
|
|
|
|
def insert_left(self, value):
|
|
self.left = Test.BinaryTreeNode(value)
|
|
return self.left
|
|
|
|
def insert_right(self, value):
|
|
self.right = Test.BinaryTreeNode(value)
|
|
return self.right
|
|
|
|
def test_full_tree(self):
|
|
tree = Test.BinaryTreeNode(5)
|
|
left = tree.insert_left(8)
|
|
right = tree.insert_right(6)
|
|
left.insert_left(1)
|
|
left.insert_right(2)
|
|
right.insert_left(3)
|
|
right.insert_right(4)
|
|
result = is_balanced(tree)
|
|
self.assertTrue(result)
|
|
|
|
def test_both_leaves_at_the_same_depth(self):
|
|
tree = Test.BinaryTreeNode(3)
|
|
left = tree.insert_left(4)
|
|
right = tree.insert_right(2)
|
|
left.insert_left(1)
|
|
right.insert_right(9)
|
|
result = is_balanced(tree)
|
|
self.assertTrue(result)
|
|
|
|
def test_leaf_heights_differ_by_one(self):
|
|
tree = Test.BinaryTreeNode(6)
|
|
left = tree.insert_left(1)
|
|
right = tree.insert_right(0)
|
|
right.insert_right(7)
|
|
result = is_balanced(tree)
|
|
self.assertTrue(result)
|
|
|
|
def test_leaf_heights_differ_by_two(self):
|
|
tree = Test.BinaryTreeNode(6)
|
|
left = tree.insert_left(1)
|
|
right = tree.insert_right(0)
|
|
right_right = right.insert_right(7)
|
|
right_right.insert_right(8)
|
|
result = is_balanced(tree)
|
|
self.assertFalse(result)
|
|
|
|
def test_three_leaves_total(self):
|
|
tree = Test.BinaryTreeNode(1)
|
|
left = tree.insert_left(5)
|
|
right = tree.insert_right(9)
|
|
right.insert_left(8)
|
|
right.insert_right(5)
|
|
result = is_balanced(tree)
|
|
self.assertTrue(result)
|
|
|
|
def test_both_subtrees_superbalanced(self):
|
|
tree = Test.BinaryTreeNode(1)
|
|
left = tree.insert_left(5)
|
|
right = tree.insert_right(9)
|
|
right_left = right.insert_left(8)
|
|
right.insert_right(5)
|
|
right_left.insert_left(7)
|
|
result = is_balanced(tree)
|
|
self.assertFalse(result)
|
|
|
|
def test_both_subtrees_superbalanced_two(self):
|
|
tree = Test.BinaryTreeNode(1)
|
|
left = tree.insert_left(2)
|
|
right = tree.insert_right(4)
|
|
left.insert_left(3)
|
|
left_right = left.insert_right(7)
|
|
left_right.insert_right(8)
|
|
right_right = right.insert_right(5)
|
|
right_right_right = right_right.insert_right(6)
|
|
right_right_right.insert_right(9)
|
|
result = is_balanced(tree)
|
|
self.assertFalse(result)
|
|
|
|
def test_only_one_node(self):
|
|
tree = Test.BinaryTreeNode(1)
|
|
result = is_balanced(tree)
|
|
self.assertTrue(result)
|
|
|
|
def test_linked_list_tree(self):
|
|
tree = Test.BinaryTreeNode(1)
|
|
right = tree.insert_right(2)
|
|
right_right = right.insert_right(3)
|
|
right_right.insert_right(4)
|
|
result = is_balanced(tree)
|
|
self.assertTrue(result)
|
|
|
|
|
|
unittest.main(verbosity=2)
|