LeetCode 110: Balanced Binary Tree

Problem Restatement

We are given the root of a binary tree.

We need to determine whether the tree is height-balanced. A binary tree is height-balanced when the left and right subtrees of every node differ in height by no more than 1. The official problem asks us to return true if the tree is balanced, otherwise false.

For this tree:

        3
      /   \
     9     20
          /  \
         15   7

The tree is balanced because every node has left and right subtree heights differing by at most 1.

For this tree:

          1
        /   \
       2     2
      / \
     3   3
    / \
   4   4

The tree is not balanced because the left side becomes too deep.

Input and Output

LeetCode gives the TreeNode class:

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

The function shape is:

class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        ...

Examples

Consider this tree:

        3
      /   \
     9     20
          /  \
         15   7

The left subtree of 3 has height 1.

The right subtree of 3 has height 2.

The difference is:

abs(1 - 2) = 1

That is allowed.

The subtrees rooted at 9, 15, and 7 are also balanced.

So the answer is:

True

Now consider:

          1
        /   \
       2     2
      / \
     3   3
    / \
   4   4

At the left child 2, the left subtree is much deeper than the right subtree.

The height difference is greater than 1.

So the answer is:

False

For an empty tree:

root = None

The answer is:

True

First Thought: Check Height at Every Node

The direct definition says:

For every node, compare the height of its left subtree and right subtree.

So we could write:

1. Compute height(root.left).
2. Compute height(root.right).
3. Check whether their difference is at most 1.
4. Recursively check the left and right subtrees.

This works, but it repeats height calculations.

For example, when checking the root, we compute the height of a subtree. Then when checking that subtree’s root, we compute many of the same heights again.

That can lead to O(n^2) time on a skewed tree.

Key Insight

We can compute height and balance at the same time.

Use a bottom-up DFS.

For each node, ask its children for their heights.

Then:

if either child subtree is already unbalanced:
    this subtree is unbalanced

if abs(left_height - right_height) > 1:
    this subtree is unbalanced

otherwise:
    this subtree is balanced, and its height is 1 + max(left_height, right_height)

A clean way to encode this is to return -1 when a subtree is unbalanced.

So the helper function returns:

Algorithm

Define a helper function:

height(node)

For each node:

1. If node is None, return 0.
2. Recursively compute the left subtree height.
3. If the left subtree returned -1, return -1.
4. Recursively compute the right subtree height.
5. If the right subtree returned -1, return -1.
6. If the height difference is greater than 1, return -1.
7. Otherwise, return 1 + max(left_height, right_height).

At the end:

return height(root) != -1

Correctness

The helper function returns the height of a subtree exactly when that subtree is balanced. If the subtree is not balanced, it returns -1.

For an empty subtree, the function returns 0, which is the correct height and is balanced.

For a non-empty node, the algorithm first checks the left and right subtrees. If either subtree is already unbalanced, then the current subtree is also unbalanced, so returning -1 is correct.

If both subtrees are balanced, the current node is balanced exactly when the difference between the two subtree heights is at most 1. If the difference is greater than 1, the algorithm returns -1.

Otherwise, the current subtree is balanced, and its height is one more than the larger child height:

1 + max(left_height, right_height)

By applying this rule from leaves upward to the root, the algorithm correctly detects whether every node satisfies the balance rule. Therefore, height(root) != -1 returns True exactly when the whole tree is height-balanced.

Complexity

Here, n is the number of nodes and h is the height of the tree.

For a balanced tree, h = O(log n).

For a skewed tree, h = O(n).

Implementation

from typing import Optional

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        def height(node: Optional[TreeNode]) -> int:
            if node is None:
                return 0

            left_height = height(node.left)
            if left_height == -1:
                return -1

            right_height = height(node.right)
            if right_height == -1:
                return -1

            if abs(left_height - right_height) > 1:
                return -1

            return 1 + max(left_height, right_height)

        return height(root) != -1

Code Explanation

The helper returns the height when the subtree is balanced:

def height(node: Optional[TreeNode]) -> int:

An empty subtree has height 0:

if node is None:
    return 0

Compute the left subtree height:

left_height = height(node.left)

If the left subtree is already unbalanced, stop early:

if left_height == -1:
    return -1

Compute the right subtree height:

right_height = height(node.right)

If the right subtree is already unbalanced, stop early:

if right_height == -1:
    return -1

Now check the current node’s balance rule:

if abs(left_height - right_height) > 1:
    return -1

If the subtree is balanced, return its height:

return 1 + max(left_height, right_height)

Finally, the full tree is balanced when the root helper call does not return -1:

return height(root) != -1

Testing

from typing import Optional

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        def height(node: Optional[TreeNode]) -> int:
            if node is None:
                return 0

            left_height = height(node.left)
            if left_height == -1:
                return -1

            right_height = height(node.right)
            if right_height == -1:
                return -1

            if abs(left_height - right_height) > 1:
                return -1

            return 1 + max(left_height, right_height)

        return height(root) != -1

def run_tests():
    s = Solution()

    root1 = TreeNode(
        3,
        TreeNode(9),
        TreeNode(20, TreeNode(15), TreeNode(7)),
    )
    assert s.isBalanced(root1) is True

    root2 = TreeNode(
        1,
        TreeNode(
            2,
            TreeNode(3, TreeNode(4), TreeNode(4)),
            TreeNode(3),
        ),
        TreeNode(2),
    )
    assert s.isBalanced(root2) is False

    root3 = None
    assert s.isBalanced(root3) is True

    root4 = TreeNode(1)
    assert s.isBalanced(root4) is True

    root5 = TreeNode(
        1,
        TreeNode(2),
        TreeNode(3, TreeNode(4), None),
    )
    assert s.isBalanced(root5) is True

    print("all tests passed")

run_tests()

Test meaning: