Balanced Binary Tree

You'll visit children before their parent (postorder), so each node's height is computed once and reused. Read Maximum Depth first. Related chapter: trees.

1. The problem

You get the root of a binary tree. Return whether, at every node, the left and right subtree heights differ by at most one. Height is the number of nodes on the longest path down, and a missing subtree has height 0. Don't change the tree. The inputs below list the tree level by level. This is LeetCode 110.

"Every node" is the phrase that matters. Checking only the root isn't enough, as section 6 shows.

2. A correct baseline

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

def height(node):
    return 0 if node is None else 1 + max(height(node.left), height(node.right))

def baseline(root):
    if root is None:
        return True
    local_ok = abs(height(root.left) - height(root.right)) <= 1
    left_ok = baseline(root.left)
    right_ok = baseline(root.right)
    return local_ok and left_ok and right_ok

assert baseline(TreeNode(2, TreeNode(1), TreeNode(3)))
assert not baseline(TreeNode(1, TreeNode(2, TreeNode(3))))
assert baseline(None)

On a chain of n nodes, each node recomputes the height of everything below it: n + (n − 1) + … + 1 visits. That's O(n²) time. On a 5,000-node chain, it's about 12 million visits. The recursion uses O(h) stack space.

The bottleneck. It computes a subtree's height, throws it away, then computes it again for every ancestor.

3. The observation

In the chain 1 → 2 → 3, node 3's height is 1, so node 2's is 2, so node 1's is 3. Each height comes straight from the children's heights. If you finish the children first, you never need to walk down again.

So visit each node's children before the node itself. That order is called postorder. With both child heights known, checking the node and computing its height take O(1) time. Store finished heights in a dictionary keyed by id(node). Using your own stack instead of recursion avoids Python's recursion limit on long chains.

4. The solution and a trace

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

class Solution:
    def isBalanced(self, root) -> bool:
        """Check every node of an acyclic binary tree; empty is balanced."""
        heights = {id(None): 0}
        pending = [(root, False)]
        while pending:
            node, expanded = pending.pop()
            if node is None:
                continue
            if not expanded:
                # Revisit this parent after both children have heights.
                pending.append((node, True))
                pending.append((node.right, False))
                pending.append((node.left, False))
            else:
                left = heights[id(node.left)]
                right = heights[id(node.right)]
                if abs(left - right) > 1:
                    return False
                heights[id(node)] = 1 + max(left, right)
        return True

solve = Solution().isBalanced
assert solve(TreeNode(2, TreeNode(1), TreeNode(3)))
assert not solve(TreeNode(1, TreeNode(2, TreeNode(3))))
assert solve(None)

Trace the left-only chain 1 → 2 → 3. Each row is a node whose children are done:

5. Why it works

What stays true. When a node is checked, both its child heights are already correct. Missing children count as 0, and real children always finish first.

A False answer is right. If the two heights differ by more than one, that node breaks the rule.

A True answer is right. Otherwise the node's saved height is correct, and checking continues up the tree. Returning True means every node passed, not just the root.

6. Cost, edge cases, and failure

Cost. Each node is handled twice, once going down and once coming back: O(n) time. The heights dictionary takes O(n) space, and the stack O(h). A recursive version would need only O(h) space, but would crash on very deep trees in Python.

Edge cases. An empty tree is balanced. A single node is balanced.

Where it fails. Equal heights at the root don't mean the whole tree is balanced:

left = TreeNode(2, TreeNode(3, TreeNode(4)))
right = TreeNode(5, None, TreeNode(6, None, TreeNode(7)))
root = TreeNode(1, left, right)
assert height(left) == height(right)
assert not solve(root)  # Both child subtrees are unbalanced.

import random
rng = random.Random(110)
def random_tree(depth):
    if depth == 0 or rng.randrange(4) == 0:
        return None
    return TreeNode(rng.randrange(4), random_tree(depth - 1), random_tree(depth - 1))
def snapshot(root):
    nodes, pending = [], [root]
    while pending:
        node = pending.pop()
        if node is not None:
            nodes.append((node, node.val, node.left, node.right))
            pending.extend((node.left, node.right))
    return nodes
for _ in range(500):
    tree = random_tree(7)
    before = snapshot(tree)
    assert solve(tree) == baseline(tree)
    assert all(node.val == val and node.left is left and node.right is right
               for node, val, left, right in before)
# Stress extension beyond the official 5,000-node limit.
chain = None
for _ in range(10000):
    chain = TreeNode(0, chain)
assert not solve(chain)