Complete balanced binary tree problem

Supporting a function that returns true if a tree has no two leaf nodes with
depth differences greater than 1.

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)