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 = Nonedef 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 casesif __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 casesif __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])