cracking the code interview balanced tree python

Implement a function to check if a binary tree is balanced.

Solution:

We need to define two functions one to check the depth of nodes, another need to compare the depth.

Firstly, we build the test case:

class TreeNode:
    def __init__(self,val):
        self.val=val
        self.left=None
        self.right=None

root=TreeNode(0)
root.left=TreeNode(1)
root.right=TreeNode(2)
root.left.left=TreeNode(3)
root.left.right=TreeNode(4)
root.left.left.left=TreeNode(5)
root.right.left=TreeNode(8)
root.right.right=TreeNode(9)

Now we have a depth function:

def depth(root):
    if root==None:
        return 0
    if root.left==None and root.right==None:
        return 1
    return 1+max(depth(root.left),depth(root.right))

Finally we combine the depth function with balance function

def balance(root):
    if root==None:
        return True
    elif abs(depth(root.left)-depth(root.right))>1:
        return False
    return balance(root.left) and balance(root.right)
    
def main():
    result=balance(root)
    print result
if __name__=="__main__":
    main()