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