superbalanced.py

import unittest

def dfs(parent, leaves,visited,depth):
    visited.append(parent)
    leftChild = parent.left
    rightChild = parent.right
    depth +=1
    if not(leftChild == None) :
        if leftChild. left == None and leftChild. right == None:
            leaves[leftChild] = depth
            visited.append(leftChild)
        else:
            dfs(leftChild,leaves,visited,depth)
    
    if not(rightChild == None) :
        if rightChild. left == None and rightChild. right == None:
            leaves[rightChild] = depth
            visited.append(rightChild)
        else:
            dfs(rightChild,leaves,visited,depth)
        
    return leaves

def collectLeaves(tree_root):
    leaves = dfs(tree_root,{},[],0)
    return leaves
def is_balanced(tree_root):

    # Determine if the tree is superbalanced
    leaves = collectLeaves(tree_root)
    if len(leaves)==0 :
        return True
    first = leaves [ list(leaves.keys())[0]]
    for key in leaves.keys():
        print (first, leaves[key])
        if not(first == leaves[key] or first == leaves[key]+1):
            return False
    print('-----------')
    return True


















# 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):
        print('test_leaf_heights_differ_by_one')
        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_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_left(2)
        right_right = right.insert_right(3)
        right_right.insert_right(4)
        result = is_balanced(tree)
        self.assertTrue(result)


unittest.main(verbosity=2)