Lesson 32 — Trees V: Piggyback on Height

~20 minutes · Days 35-36 · tree DFS return values · Python

Tonight you meet two tree interview problems that look bigger than they are: Diameter of Binary Tree and Balanced Binary Tree. Both ask for a whole-tree fact, but the fact is hiding inside the same small question at every node: how tall are my children?

The invariant

Compute height recursively and carry the real answer up alongside it — one traversal answers both questions.

That's the pattern. In tree DFS, the answer for a node usually comes from the answers of its children. Here the child answer you must return is height, and the extra answer you carry is the thing the problem really asked for. Python can return both as one tuple; the language docs call these tuples and sequences.

Diameter: the slow way first

The diameter is the longest path between any two nodes. The brute force idea is: for every node, compute the height of its left subtree and right subtree from scratch, then try leftHeight + rightHeight. That works, but on a skinny tree it can become O(n²), because the same lower nodes get re-counted again and again.

The fix is to make each DFS call do two jobs: return its height to the parent, and carry the best diameter found inside that subtree.

The node sends height upward, but it also checks whether the best diameter passes through its left and right subtrees.

def diameterOfBinaryTree(root):
    def dfs(node):
        if not node:
            return 0, 0              # height, best diameter
        lh, ld = dfs(node.left)      # child height + child answer
        rh, rd = dfs(node.right)
        height = 1 + max(lh, rh)     # what parent needs
        best = max(ld, rd, lh + rh)  # what this subtree knows
        return height, best
    return dfs(root)[1]

Notice the split: height is the value the parent needs, while best is the answer the whole problem needs. A node is useful because it can join the best downward path from the left child with the best downward path from the right child.

Transfer: same pattern, new question

Balanced Binary Tree: return whether every node has left and right subtree heights that differ by at most one.

Before opening the reveal: predict the two things each DFS call should return. One is still height. What is the carried answer now?

def isBalanced(root):
    def dfs(node):
        if not node:
            return 0, True           # height, is balanced
        lh, lb = dfs(node.left)
        rh, rb = dfs(node.right)
        height = 1 + max(lh, rh)
        ok = lb and rb and abs(lh - rh) <= 1
        return height, ok
    return dfs(root)[1]

What the trade buys you

The space is recursion stack height h. The win is time: every node computes its height once, then hands both useful facts upward.

Practice — answer before you scroll past

lh, ok = dfs(node.left)
rh, ok = dfs(node.right)
return 1 + max(lh, rh), ok and abs(lh - rh) <= 1

Your move (tonight)

Solve these on LeetCode, in order, logging each in your Notion tracker with the cycle — predict → attempt → compare → encode:

1. Diameter of Binary Tree — predict the brute force first, then write the tuple-return DFS from memory.
2. Balanced Binary Tree — use the same skeleton, changing only the carried answer.

For the encode step, write the invariant at the top of this page in your own words. The tree self-check gate is: write a DFS-with-return-value template cold, then explain diameter piggybacking.