Chapter 4 Trees and Graphs - 4.1

Problem 4.1: Implement a function to check if a tree is balanced. For the purposes of this question, a balanced tree is defined to be a tree such that no two leaf nodes differ in distance from the root by more than one.

The answer page gives a solution that is very simple to implement:

class binary_tree_node:
    def __init__(self, value = None):
        self.value = value
        self.left_child = None
        self.right_child = None

def is_balanced(root):
    return (maxDepth(root) - minDepth(root) <= 1)
def maxDepth(node):
    if node == None:
        return 0
    return (1 + max([maxDepth(node.left_child), maxDepth(node.right_child)]))
def minDepth(node):
    if node == None:
        return 0
    return (1 + min([minDepth(node.left_child), minDepth(node.right_child)]))
# Test cases
if __name__ == "__main__":
    nodes = []
    for i in range(0, 10):
        nodes.append(binary_tree_node(i))
    nodes[0].left_child = nodes[1]
    nodes[0].right_child = nodes[2]
    nodes[1].left_child = nodes[3]
    nodes[1].right_child = nodes[4]
    nodes[4].left_child = nodes[5]
    print is_balanced(nodes[0])

I have to admit that the readability of solution above is quite high. However, there are two points can be improved: 1) the algorithm should be applied to any kind of trees (not only to binary tree), according to the description of problem; 2) traversing the whole tree is not necessary in some cases and we can actually stop traversing the tree if we find it is not balanced.

My solution is given below:

from queue import *
class tree_node:
    def __init__(self, value = None):
        self.value = value
        self.children = []
def is_balanced(root):
    # Use BFS
    q = queue()
    q.enqueue(root)
    current_level_num = 1
    next_level_num = 0
    current_depth = 0
    is_first_leaf = True
    min_leaf_depth = 0
    while not q.is_empty():
        node = q.dequeue()
        current_level_num -= 1
        # If current node is the first leaf in BFS
        # its depth is the minimum
        if is_first_leaf and (len(node.children) == 0):
            min_leaf_depth = current_depth
            is_first_leaf = False
        # If current node is a leaf and it is not the first node
        # in BFS, we compare its depth to minimum one
        if (not is_first_leaf) and (len(node.children) == 0):
            if (current_depth - min_leaf_depth) > 1:
                return False
        for n in node.children:
            next_level_num += 1
            q.enqueue(n)
        # After we traverse a level
        if current_level_num == 0:
            current_level_num = next_level_num
            next_level_num = 0
            current_depth += 1
    return True
# Test cases
if __name__ == "__main__":
    nodes = []
    for i in range(0, 10):
        nodes.append(tree_node(i))
    nodes[0].children.append(nodes[1])
    nodes[0].children.append(nodes[2])
    nodes[0].children.append(nodes[3])
    nodes[1].children.append(nodes[4])
    nodes[1].children.append(nodes[5])
    nodes[1].children.append(nodes[6])
    #nodes[4].children.append(nodes[7])
    print is_balanced(nodes[0])