Implement a function to check if a tree is balanced.

is_tree_balanced.py

# -*- coding: utf-8 -*-
"""
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.
"""


class Node(object):
    def __init__(self, value):
      self.left = None
      self.right = None
      self.value = value


def create_balanced_tree():
    """Generates the following tree:

       (1)
       / \
     (2)  (3)
     /  \   \
    (4) (5)  (6)
          \
          (7)
    """
    root = Node(1)
    root.left = Node(2)
    root.right = Node(3)
    root.right.right = Node(6)
    root.left.left = Node(4)
    root.left.right = Node(5)
    root.left.right.left = Node(7)
    return root


def create_unbalanced_tree():
    """Generates the following tree:

       (1)
       / \
     (2)  (3)
     /  \   \
    (4) (5)  (6)
          \
          (7)
            \
            (8)
    """
    root = Node(1)
    root.left = Node(2)
    root.right = Node(3)
    root.right.right = Node(6)
    root.left.left = Node(4)
    root.left.right = Node(5)
    root.left.right.left = Node(7)
    root.left.right.left.right = Node(8)
    return root


def is_balanced(t):
    if t is None:
        return 0
    left_height = is_balanced(t.left)
    if left_height == -1:
        return -1
    right_height = is_balanced(t.right)
    if right_height == -1:
        return -1
    diff = left_height - right_height
    if abs(diff) > 1:
        return -1
    return 1 + max(left_height, right_height)


def check_is_balanced(t):
    if is_balanced(t) < 1:
        return False
    return True


if __name__ == '__main__':
    tree1 = create_balanced_tree()
    tree2 = create_unbalanced_tree()
    print "tree1: {0}".format(check_is_balanced(tree1))
    print "tree2: {0}".format(check_is_balanced(tree2))