about summary refs log tree commit diff
path: root/users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py
diff options
context:
space:
mode:
authorVincent Ambo <mail@tazj.in>2021-12-13T22·51+0300
committerVincent Ambo <mail@tazj.in>2021-12-13T23·15+0300
commit019f8fd2113df4c5247c3969c60fd4f0e08f91f7 (patch)
tree76a857f61aa88f62a30e854651e8439db77fd0ea /users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py
parent464bbcb15c09813172c79820bcf526bb10cf4208 (diff)
parent6123e976928ca3d8d93f0b2006b10b5f659eb74d (diff)
subtree(users/wpcarro): docking briefcase at '24f5a642' r/3226
git-subtree-dir: users/wpcarro
git-subtree-mainline: 464bbcb15c09813172c79820bcf526bb10cf4208
git-subtree-split: 24f5a642af3aa1627bbff977f0a101907a02c69f
Change-Id: I6105b3762b79126b3488359c95978cadb3efa789
Diffstat (limited to 'users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py')
-rw-r--r--users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py126
1 files changed, 126 insertions, 0 deletions
diff --git a/users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py b/users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py
new file mode 100644
index 000000000000..03de0350d898
--- /dev/null
+++ b/users/wpcarro/scratch/deepmind/part_two/balanced-binary-tree.py
@@ -0,0 +1,126 @@
+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)