Maximum Depth of Binary Tree

You'll walk the tree with DFS, using your own stack, and carry each node's depth with it. You need tree nodes and a stack. Related chapter: trees.

1. The problem

You get the root of a binary tree. Return the most nodes on any path from the root down to a leaf. Count nodes, not edges, so an empty tree has depth 0. The values don't matter. The inputs below list the tree level by level, with null for a missing child. This is LeetCode 104.

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 example():
    return TreeNode(3,TreeNode(9),TreeNode(20,TreeNode(15),TreeNode(7)))

def baseline(root):
    if root is None:
        return 0
    paths=[]
    pending=[(root,[])]
    while pending:
        node,path=pending.pop()
        path=path+[node.val]
        if node.left is None and node.right is None:
            paths.append(path)
        if node.left: pending.append((node.left,path))
        if node.right: pending.append((node.right,path))
    return max(map(len,paths))

assert baseline(example())==3
assert baseline(TreeNode(8))==1
assert baseline(None)==0

This builds a full copy of every path. For n nodes and height h, that's O(nh) time and space in the worst case. On a chain of 10,000 nodes, that's about 50 million values copied.

The bottleneck. It stores every value on each path, but only the path's length matters.

3. The observation

In the example, node 20 is at depth 2, so its children 15 and 7 are at depth 3. You don't need the path [3, 20] to know that; you only need the number 2.

So store each node with its depth: a pair like (node, 2). A child's depth is its parent's plus one. Keep the largest depth seen. Using your own stack, instead of recursion, also avoids Python's recursion limit on a very deep tree.

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 maxDepth(self, root) -> int:
        """Count nodes on the longest path in an acyclic binary tree."""
        if root is None:
            return 0
        pending=[(root,1)]  # Each depth is measured from the original root.
        best=0
        while pending:
            node,depth=pending.pop()
            best=max(best,depth)
            if node.left: pending.append((node.left,depth+1))
            if node.right: pending.append((node.right,depth+1))
        return best

solve=Solution().maxDepth
assert solve(example())==3
assert solve(TreeNode(8))==1
assert solve(None)==0

Trace the example tree:

5. Why it works

What stays true. Every depth on the stack is the true number of nodes from the root to that node. The root starts at 1, and each child adds one.

The answer is right. In a tree, every node is pushed exactly once, so we see every depth. The deepest node must be a leaf, because any child would be deeper. So the largest depth seen is the longest root-to-leaf path.

6. Cost, edge cases, and failure

Cost. Each node is pushed and popped once: O(n) time, growing in step with the tree. The stack holds O(h) entries, since each node on the current path leaves at most one sibling waiting. A recursive version has the same cost, but a one-sided tree of 10,000 nodes would raise RecursionError.

Edge cases. An empty tree returns 0. A single node returns 1.

Where it fails. Counting edges instead of nodes gives an answer one too small:

assert solve(TreeNode(1))==1  # Counting edges would incorrectly give 0.
assert solve(TreeNode(1,None,TreeNode(2)))==2

import random
rng=random.Random(104)
def random_tree(depth):
    if depth==0 or rng.randrange(4)==0: return None
    return TreeNode(rng.randrange(10),random_tree(depth-1),random_tree(depth-1))
for _ in range(500):
    root=random_tree(7)
    assert solve(root)==baseline(root)
chain=None
for _ in range(10000): chain=TreeNode(1,chain)
assert solve(chain)==10000