Skip to content

Instantly share code, notes, and snippets.

@nficano
Created May 19, 2016 14:44
Show Gist options
  • Save nficano/34f0dc4b8a89604dc9a5c5a230f3310f to your computer and use it in GitHub Desktop.
Save nficano/34f0dc4b8a89604dc9a5c5a230f3310f to your computer and use it in GitHub Desktop.
Implement a function to check if a tree is balanced.
# -*- 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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment