47a11b76a2
Write a predicate for determining if a binary tree is "super balanced", which means that the depths of all of the tree's leaves are equal or differ by at most one.
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)
|