Close

Checking if a binary tree is balanced

Given a binary tree, how do we check if it is balanced or not?
We say that a binary tree is balanced if height difference of left  and the right sub-trees is at most 1 at  any node.
Consider the following examples
  1          1               1             1
 /                       /              
2   3          2           2                 2
                                           /
                                          3
Balanced    Balanced     Balanced       Unbalanced
Solution:
We can simply implement a recursive solution based on the above definition of balance. At each node we check the height of the left sub tree and the right sub-tree and call the same function for it’s left child and right child recursively. But this is inefficient because we calculate the height repeatedly. This solution is O(n2).
Because we are calculating the height at each node, we can reuse the height of sub-tree to check the balance at the root node. Here is how we can do it. 
We write a method called balance. This method returns the actual height of the tree if it is balanced. Otherwise it returns -1. So at any point if the balance of left or right sub-tree is -1 or their difference is greater than 1, we can simply return false.
Here is the C++ implementation of this algorithm. It runs in O(n) time.

#include <iostream>
#include <stack>
using namespace std;
//Binary tree node type
struct TreeNode
{
int val;
TreeNode *left, *right;
TreeNode(int v):val(v),left(NULL), right(NULL)
{
}
};
//Binary search tree insert procedure for creating binary tree to test
TreeNode* insert(TreeNode * root, int data)
{
if( root == NULL )
{
root = new TreeNode(data);
return root;
}
if( data < root->val )
root->left = insert( root->left, data );
else
root->right = insert( root->right, data );
return root;
}
int height(TreeNode *root)
{
if( root == NULL )
return 0;
return max(height(root->left), height(root->right))+1;
}
//Brute force method based on recursion
bool is_balanced_bruteforce(TreeNode *root)
{
if( root == NULL )
return true;
if( abs( height(root->left)-height(root->right)) > 1 )
return false;
return is_balanced_bruteforce(root->left) && is_balanced_bruteforce(root->right);
}
//This method returns -1 if the tree is unbalanced, otherwise it returns the height of the tree
int balance(TreeNode *root)
{
if( root == NULL )
return 0;
int left_balance = balance(root->left);
if(left_balance == -1)
return -1;
int right_balance = balance(root->right);
if( right_balance == -1)
return -1;
if( abs(left_balance-right_balance) > 1 )
return -1;
return max(left_balance, right_balance)+1;
}
bool is_balanced(TreeNode *root)
{
if( balance(root) == -1)
return false;
return true; //If the balance is not -1; it is a balanced tree
}
void Test1()
{
TreeNode *root = NULL;
root = insert(root, 2);
root = insert(root, 1);
root = insert(root, 4);
if( is_balanced(root) )
cout <<"Balanced" << endl;
else
cout << "Not balanced" << endl;
}
void Test2()
{
TreeNode *root = NULL;
root = insert(root, 4);
root = insert(root, 2);
root = insert(root, 1);
root = insert(root, 6);
root = insert(root, 8);
root = insert(root, 7);
root = insert(root, 9);
if( is_balanced(root) )
cout <<"Balanced" << endl;
else
cout << "Not balanced" << endl;
}
int main()
{
Test1();
Test2();
return 0;
}

Leave a Reply

Your email address will not be published. Required fields are marked *