Finding if the tree is balanced or not

Question

Can someone please review my code and let me know if there are some bugs or possible improvements?

    /**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    //return true if the root node is null or there is only root node.
    public boolean isBalanced(TreeNode root) {
        if (root == null || (root != null && root.left == null && root.right == null)) {
            return true;
        } else {
            int balancedRight = 0;
            int balancedLeft = 0;
            if (root.left != null) {
                balancedLeft = findHeight(root.left);
            }
            if (root.right != null) {
                balancedRight = findHeight(root.right);
            }

            return Math.abs(balancedLeft - balancedRight) <= 1 ? true : false;
        }

    }
//find the height of the tree
    public int findHeight(TreeNode root) {
        if (root == null) {
            return 0;
        } else if (root.left == null && root.right == null) {
            return 1;
        } else if (root.left == null && root.right != null) {
            return 1 + findHeight(root.right);
        } else if (root.left != null && root.right == null) {
            return 1 + findHeight(root.left);
        } else {
            return Math.max(1+findHeight(root.left), 1+findHeight(root.right));
        }
    }
}

Answer 1

1. First of all, your algorithm checks if the root of the tree is balanced or not. Rather, it should check for 3 conditions, whether the root , left subtree and right subtree are balanced or not.

So check for this recursively,

Math.abs(leftSubtreeHeight - rightSubtreeHeight) <= 1) &&
            isBalanced(root.left) && isBalanced(root.right)

2. Moreover, your height method is doing too many unnecessary checks. It can be simply reduced to :-

   private int treeHeight(TreeNode root) {
   if(root == null)
       return 0;
   
   return Integer.max(treeHeight(root.left),treeHeight(root.right) ) +1;
   
   }
   

take this as an example:-

else if (root.left == null && root.right != null) {
        return 1 + findHeight(root.right);

in this, if root.left is null, then it will return 0. Hence, from the recursive call of the root, you'll get

return 1 + Math.max(findHeight(root.left), findHeight(root.right));

findHeight(root.left) will return 0 and you'll be left with just 1 + findHeight(root.right) which is what you have manually written.

3. Also, it's better to keep the treeHeight(TreeNode root) function private if nobody outside this class is going to use it.

4. You need not write else after this

   if (root == null || (root != null && root.left == null && root.right == null)) {
           return true;
       } else {
   

because if you returned from the if statement then anyways you won't go further. And, if the if condition doesn't return true, you'll go to the else part anyways.

5. Also, change the names from:- balancedRight to rightSubtreeHeight, balancedLeft to leftSubtreeHeight, isBalanced to isTreeBalanced

6. Additionally, you need not have this condition

   || (root != null && root.left == null && root.right == null)
   

because, if the root.left == null, it's height would be 0 and 0 will be returned from the right hand side as well. So, for root node, you received 0 from both left and right. Now, the difference between them is 0 which is <=1. The above code is anyways checking that using recursion.

Try to make the entire recursion tree for various problems to have a better understanding of how recursion works. By knowing the power of recursion, you can make your code more elegant.

The whole code could be reduced to

public boolean isTreeBalanced(TreeNode root){
    if(root == null)
        return true;

    int leftSubtreeHeight = treeHeight(root.left);
    int rightSubtreeHeight = treeHeight(root.right);

    if((Math.abs(leftSubtreeHeight - rightSubtreeHeight) <= 1) &&
            isTreeBalanced(root.left) && isTreeBalanced(root.right))
        return true;

    return false;
}


private int treeHeight(TreeNode root) {
    if(root == null)
        return 0;

    return Integer.max(treeHeight(root.left),treeHeight(root.right) ) +1;

}

Edit:

However, if we look further deeply, we can find out the time complexity of the algorithm can be reduced further(from O(N^2) to O(N)) because isBalancedTree and treeHeight are going through same pattern of recursion. We can get the result in one traversal of the tree. Here is the code:-

public boolean isBalanced(TreeNode root) {
    if(root == null)
        return true;
    if(helper(root) != -1)
        return true;
    else
        return false;
}

private static int helper(TreeNode root) {
    if(root == null)
        return 0;

    int lefth = helper(root.left);
    int righth = helper(root.right);
    if(lefth == -1 || righth == -1)
        return -1;

    if( Math.abs(lefth - righth) <= 1)
        return Math.max(lefth, righth) + 1;
    else
        return -1;
}

Edit: The code could be further compressed by removing unnecessary if else statements.

    public boolean isBalanced(TreeNode root) {
    if(root == null)
        return true;
    return helper(root) != -1;
}

private static int helper(TreeNode root) {
    if(root == null)
        return 0;

    int lefth = helper(root.left);
    int righth = helper(root.right);
    if(lefth == -1 || righth == -1)
        return -1;

    if( Math.abs(lefth - righth) <= 1)
        return Math.max(lefth, righth) + 1;
    else
        return -1;
}

Answer 2

Definition of balanced

The usual definition of a balanced binary tree:

1. The left and right subtrees' heights differ by at most one, AND
2. The left subtree is balanced, AND
3. The right subtree is balanced

Your algorithm only checks the level of the root, comparing the height of the left and right branches of the root. It doesn't go deeper, so the left or right subtrees may be unbalanced.

Unnecessary conditions

This condition can be simplified:

if (root == null || (root != null && root.left == null && root.right == null)) {

You can drop the root != null, because it will always be true. This is exactly the same:

if (root == null || (root.left == null && root.right == null)) {

This chain of conditions has the same issue, with some unnecessary conditions that can be eliminated:

if (root == null) {
    return 0;
} else if (root.left == null && root.right == null) {
    return 1;
} else if (root.left == null && root.right != null) {
    return 1 + findHeight(root.right);
} else if (root.left != null && root.right == null) {
    return 1 + findHeight(root.left);
} else {

Can be simplified to:

if (root == null) {
    return 0;
} else if (root.left == null && root.right == null) {
    return 1;
} else if (root.left == null) {
    return 1 + findHeight(root.right);
} else if (root.right == null) {
    return 1 + findHeight(root.left);
} else {

Using boolean expressions

You can use boolean expressions directly. For example like this:

return Math.abs(balancedLeft - balancedRight) <= 1;

Other simplifications

Instead of this:

int balancedRight = 0;
int balancedLeft = 0;
if (root.left != null) {
    balancedLeft = findHeight(root.left);
}
if (root.right != null) {
    balancedRight = findHeight(root.right);
}

Since findHeight immediately does a check for the case of root == null, you could simplify the above as:

int balancedRight = findHeight(root.left);
int balancedLeft = findHeight(root.right);

Another thing, instead of adding a constant in both sides of a Math.max:

return Math.max(1+findHeight(root.left), 1+findHeight(root.right));

You could do just once:

return 1 + Math.max(findHeight(root.left), findHeight(root.right));

On closer look, you could simplify the entire findHeight method to this:

public int findHeight(TreeNode root) {
    if (root == null) {
        return 0;
    }
    return 1 + Math.max(findHeight(root.left), findHeight(root.right));
}

Naming

Instead balancedLeft and balancedRight, leftHeight and rightHeight would have made a lot more sense.