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:
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])
