diff --git a/deepmind/balanced-binary-tree.py b/deepmind/balanced-binary-tree.py new file mode 100644 index 000000000..7fc174a2a --- /dev/null +++ b/deepmind/balanced-binary-tree.py @@ -0,0 +1,123 @@ +import unittest +from collections import deque + + +def is_balanced(node): + q, seen, ds = deque(), set(), set() + q.append((0, node)) + while q: + d, node = q.popleft() + l, r = node.left, node.right + seen.add(node) + if not l and not r: + if d not in ds and len(ds) == 2: + return False + else: + ds.add(d) + if l and l not in seen: + q.append((d + 1, l)) + if r and r not in seen: + q.append((d + 1, r)) + return max(ds) - min(ds) <= 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)