Wrong output when checking whether binary tree is balanced

Question

I am trying to solve LeetCode problem 110. Balanced Binary Tree:

Given a binary tree, determine if it is height-balanced.

Here is my attempt:

class Solution {
    boolean c = true;

    public boolean isBalanced(TreeNode root) {
        int diff = helper(root, 0);
        System.out.println(diff);
        if ((diff > 0 && diff < 2) || root == null)
            return true;
        return false;
    }

    int helper(TreeNode root, int count) {
        if (root == null) {
            return count;
        }
        int a = helper(root.left, count + 1);
        int b = helper(root.right, count + 1);
        int diff = Math.abs(a - b);
        if (diff > 1 || diff < 0)
            return count;
        return Math.max(a, b);
    }
}

I get the wrong result for the following input:

          3
         / \
        9   20
           /  \
         15    7

Encoded in level order: [3,9,20,null,null,15,7]

The expected output is true, but my code returns false.

The TreeNode class is as follows:

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode() {}
    TreeNode(int val) { this.val = val; }
    TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

Here is the main code to reproduce it:

    public static void main(String[] args) {
        Solution solution = new Solution();

        TreeNode tree = new TreeNode(3, 
            new TreeNode(9), 
            new TreeNode(20,
                new TreeNode(15),
                new TreeNode(7)
            )
        );
        
        boolean result = solution.isBalanced(tree);
        System.out.println("result " + result); // false, but should be true
    }

What is my mistake?

Answer 1

The problem with your code is that you sometimes expect helper to return a height and other times to return a difference in heights. So for instance:

• In isBalanced you expect helper(root,0) to return a difference
• In helper you expect helper(root.left,count+1) and helper(root.right,count+1) to return heights.
• In helper all return statements return a height, not a difference in heights.

You will need to make a clear distinction between the two. As your heights are all non-signed integers, you could reserve the value -1 to indicate that an imbalance was found, and if so, bubble this -1 up all the way to the initial call of helper so that isBalanced can know whether the whole tree is balanced or not by comparing the returned value with -1.

Not a problem, but your condition diff < 0 will never be true, as diff is an absolute value.

Here is a correction of your code with that idea:

class Solution {
    public boolean isBalanced(TreeNode root) {
        int result = helper(root, 0); // returns a height, or -1 if not balanced
        return result != -1;
    }

    int helper(TreeNode root,int count){
        if (root==null) {
            return count;
        }
        int a=helper(root.left,count+1);
        if (a == -1) return -1; // Don't bother continuing: we already know the tree is not balanced!
        int b=helper(root.right,count+1);
        if (b == -1) return -1; // ..idem..
        int diff=Math.abs(a-b);
        if (diff > 1) // Not balanced!
            return -1;  // Special value to indicate the tree is not balanced
        return Math.max(a, b); // The height of a balanced tree
    }
}