I am learning about the binary tree data structure and implemented some basic functionality.
It was working great with module level functions until I realized that I need to keep track of the root for operations like checking height. The simplest solution that I could possibly think was of creating a container class to hold the root reference, but all my functions became quite ugly then.
class TreeNode(object):
"""Node of a Binary tree."""
def __init__(self, key, left=None, right=None):
self.key = key
self.left = left
self.right = right
def __str__(self):
return "{0}".format(self.key)
class BinaryTree(object):
"""ADT for binary tree."""
def __init__(self, root=None):
self.root = root
def is_root(self, node):
return node.key == self.root.key
def is_leaf(self, node):
"""Checks whether given node is leaf."""
if node is None or self.is_root(node):
return False
return node.left is None and node.right is None
def is_internal(self, node):
"""Checks whether the given node is internal node.
All nodes except leaf nodes are considered internal nodes.
"""
return not self.is_leaf(node)
def is_full(self):
"""Checks if the subtree rooted at the given node is full.
A Binary tree is full when every node other than leaves
has two children.
"""
def recurse(node):
if node is None:
return True
if self.is_leaf(node):
return True
if node.left is not None and node.right is not None:
return recurse(node.left) and recurse(node.right)
return False
return recurse(self.root)
def height(self):
"""Calculates height of the Binary tree."""
return self.node_height(self.root)
def node_height(self, node):
"""Calculates height of the subtree rooted at given node."""
if node is None or self.is_leaf(node):
return 0
return max(self.node_height(node.left), self.node_height(node.right)) + 1
def is_balanced(self):
"""Checks whether the subtree rooted at the given node is balanced.
The subtree rooted at node is balanced if the difference between the
height of the left left subtree and right subtree is within [-1, 0, 1].
"""
def recurse(node):
if node is None:
return True
lheight = self.node_height(node.left)
rheight = self.node_height(node.right)
return abs(lheight - rheight) < 1
return recurse(self.root)
@classmethod
def level_order(cls, node):
"""Given Binary Tree node gives list of nodes in each level."""
if node is None:
return
current_level = [node]
res = []
while current_level:
res.append(current_level)
next_level = []
for node in current_level:
if node.left:
next_level.append(node.left)
if node.right:
next_level.append(node.right)
current_level = next_level
return res
@classmethod
def build_tree(cls, inorder, preorder):
"""Builds a Binary Tree from given inorder and preorder traversal.
Steps:
1. Take the first element from the preorder list.
2. Get the index of the element from inorder list.
3. Set the element as root.
4. Take all elements left to root from inorder list.
5. Take all elements right to root from inorder list.
6. Calculate preorder list for left and right subtree respectively.
7. Repeat from step 1 for each subtree.
"""
def insert_left(parent, node):
parent.left = node
return parent
def insert_right(parent, node):
parent.right = node
return parent
def recurse(node, inorder, preorder):
if len(preorder) == 1:
return TreeNode(preorder)
key = preorder[0]
root_index = inorder.find(key)
inorder_left = inorder[:root_index]
inorder_right = inorder[root_index+1:]
preorder_left = preorder[1:len(inorder_left)+1]
preorder_right = preorder[root_index+1:]
node = TreeNode(key)
if len(inorder_left):
insert_left(node, recurse(None, inorder_left, preorder_left))
if len(inorder_right):
insert_right(node, recurse(None, inorder_right, preorder_right))
return node
node = recurse(None, inorder, preorder)
return cls(node)
def print_tree(self, node):
node = self.root if node is None else node
for level in self.level_order(root):
for node in level:
print node,
print
import unittest
class TestBinaryTree(unittest.TestCase):
#@unittest.skip("demonstrating skipping")
def test_inorder(self):
bt = BinaryTree.build_tree('DBEAFIC', 'ABDEICF')
expected = [['A'], ['B', 'I'], ['D', 'E', 'C', 'F']]
actual = []
for level in bt.level_order(bt.root):
actual.append([node.key for node in level ])
self.assertEqual(actual, expected)
def test_isFull(self):
bt = BinaryTree.build_tree('DBEAFIC', 'ABDEICF')
self.assertTrue(bt.is_full())
bt = BinaryTree.build_tree('DBEAFC', 'ABDECF')
self.assertFalse(bt.is_full())
class TestBinaryTreeHeight(unittest.TestCase):
def test_none(self):
bt = BinaryTree()
self.assertEqual(bt.height(), 0)
def test_only_root(self):
bt = BinaryTree.build_tree('A', 'A')
self.assertEqual(bt.height(), 1)
def test_tree(self):
bt = BinaryTree.build_tree('DBEAFC', 'ABDECF')
self.assertEqual(bt.height(), 2)
class TestBinaryTreeBalance(unittest.TestCase):
def test_leftSkewed(self):
"""
A
/
B
/
C
"""
bt = BinaryTree.build_tree('CBA', 'ABC')
self.assertFalse(bt.is_balanced())
def test_rightSkewed(self):
"""
A
\
B
\
C
"""
bt = BinaryTree.build_tree('ABC', 'ABC')
self.assertFalse(bt.is_balanced())
def test_balanced(self):
"""
A
/ \
B C
/ \
D E
"""
bt = BinaryTree.build_tree('DBEAC', 'ABDEC')
self.assertFalse(bt.is_balanced())
def test_unbalanced(self):
"""
A
/ \
B C
/ \
D E
/ \
F G
"""
bt = BinaryTree.build_tree('DBFEGAC', 'ABDEFGC')
self.assertFalse(bt.is_balanced())
if __name__ == '__main__':
unittest.main()
Note
Some comments in the methods may seem arbitrary because initially they were module level functions which expected any node.
Doubts
- What strategy I could have adopted to keep my code clean?
- I didn't store a height for each node to keep my objects lean. Instead, in
is_balanced, I am calculating expensive calculations. Is this fine? - [TODO].
