Diameter of a Binary Tree: One DFS That Returns Depth

The diameter of a binary tree is the length of the longest path between any two nodes. It is a short problem with a lesson that applies to half of all tree questions: the value a recursive function returns to its parent does not have to be the answer. Here the function returns a depth and quietly records the answer on the side, and seeing why is the whole point.

The problem it solves

Given the root of a binary tree, return the number of edges on the longest path between any two nodes. The path does not have to pass through the root. In the tree with root 1, children 2 and 3, and 2's children 4 and 5, the diameter is 3: the path 4, 2, 1, 3.

The first idea that comes to mind: for every node, compute the height of its left and right subtrees and add them. That is correct, but computing heights from scratch at every node visits each subtree again and again, and on a tree shaped like a line it costs O(n^2).

The intuition

Every path in a tree has exactly one highest node, the place where it bends: it comes up one side and goes down the other, or it starts at that node and goes down one side. The longest path that bends at node n goes as deep as possible on the left and as deep as possible on the right, so its length is left depth + right depth, counting depth as the number of nodes on the longest downward path from a child.

So one post-order traversal does everything:

• Ask the left child for its depth and the right child for its depth.
• Record left + right as a candidate: the best path that bends here.
• Return 1 + max(left, right) to the parent: the parent can only continue a path down one side, never both.

The return value and the recorded value are different things. The parent needs a depth, because a path through it can use only one branch of this node. The global answer needs the bend, because a path that bends here uses both. Confusing the two is the most common bug.

Watch it run

The animation walks the lesson's tree bottom-up and prints each node's returned depth under it and its candidate bend above it. Leaves 4 and 5 each report a depth of 1. At 2, bending costs 1 + 1 = 2 and it reports 2. Leaf 3 reports 1. At 1, bending costs 2 + 1 = 3 and it reports 3. The last frame highlights the winning path, three edges from leaf 4 to leaf 3, found in one visit per node.

The code

The recursive version, a tree where the longest path avoids the root, and the empty and one-node cases:

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

def diameter(root):
    best = 0
    def depth(node):                           # nodes on the longest downward path
        nonlocal best
        if node is None:
            return 0
        l, r = depth(node.left), depth(node.right)
        best = max(best, l + r)                # edges on the path that bends here
        return 1 + max(l, r)                   # what this node reports upwards
    depth(root)
    return best

root = Node(1, Node(2, Node(4), Node(5)), Node(3))
print(diameter(root))                          # 3  (4-2-1-3)

lopsided = Node(1, Node(2, Node(3, Node(4), None), Node(5, None, Node(6, None, Node(7)))))
print(diameter(lopsided))                      # 5  (4-3-2-5-6-7, root not on it)
print(diameter(None), diameter(Node(9)))       # 0 0

Recursion depth equals tree height, and Python's default limit is about a thousand frames. A degenerate tree of 5,000 nodes in a line needs an explicit stack, doing the same post-order work:

def diameter_iterative(root):
    if root is None:
        return 0
    best, height, stack = 0, {None: 0}, [(root, False)]
    while stack:
        node, children_done = stack.pop()
        if children_done:                      # post-order: both children measured
            l, r = height[node.left], height[node.right]
            best = max(best, l + r)
            height[node] = 1 + max(l, r)
        else:
            stack.append((node, True))
            for child in (node.left, node.right):
                if child:
                    stack.append((child, False))
    return best

chain = None
for v in range(5000):                          # 5,000 nodes in a line
    chain = Node(v, chain)
print(diameter_iterative(chain))               # 4999
try:
    diameter(chain)
except RecursionError:
    print("RecursionError")                    # RecursionError

Both against a brute force that runs a BFS from every node over the tree's undirected edges and keeps the longest shortest path, on 300 random trees:

import random
from collections import deque

def random_tree(n):
    nodes = [Node(i) for i in range(n)]
    for i in range(1, n):
        while True:
            parent = nodes[random.randrange(i)]
            side = random.choice(["left", "right"])
            if getattr(parent, side) is None:
                setattr(parent, side, nodes[i])
                break
    return nodes

def brute(nodes):
    """BFS from every node over the undirected edges; longest shortest path."""
    adj = {n: [] for n in nodes}
    for n in nodes:
        for child in (n.left, n.right):
            if child:
                adj[n].append(child)
                adj[child].append(n)
    best = 0
    for start in nodes:
        dist, q = {start: 0}, deque([start])
        while q:
            u = q.popleft()
            for v in adj[u]:
                if v not in dist:
                    dist[v] = dist[u] + 1
                    q.append(v)
        best = max(best, max(dist.values()))
    return best

random.seed(16)
ok = True
for _ in range(300):
    nodes = random_tree(random.randint(1, 25))
    ok &= diameter(nodes[0]) == diameter_iterative(nodes[0]) == brute(nodes)
print(ok)                                      # True

The complexity

• Heights recomputed at every node: O(n^2) in the worst case, a tree shaped like a line.
• One post-order DFS: each node is visited once and does constant work, so O(n) time. Space is O(h) for the recursion or the explicit stack, where h is the height: O(log n) for a balanced tree, O(n) for a line.
• The brute force: a BFS from each of n nodes, O(n^2); fine as a test, not as an answer.

Where it goes wrong

• Returning the diameter instead of the depth. A parent that receives a bend length will add both of a child's branches to its own path, which is not a path at all.
• Only checking the root. left depth + right depth at the root misses the lopsided case above, where the answer is 5 but the root only sees 4.
• Edges versus nodes. This problem counts edges. If a version asks for nodes on the path, it is best + 1. Say which one you are returning.
• Deep trees in Python. Recursion hits the default limit on a long chain; mention the iterative version or raising the limit.

When it shows up in interviews

It shows up as an easy-to-medium tree question, and it is often the warm-up for harder problems with the same "return one thing, record another" shape: binary tree maximum path sum, where the bend adds node values and a branch may be dropped if it is negative, and longest univalue path. For a general tree or graph that is a tree, a second classic method runs BFS from any node to the farthest node, then BFS again from there; the second distance is the diameter. In practice the same idea measures the widest spread of a network, like the longest hop count between two routers in a tree-shaped topology.

How to say it in an interview

"Every path has one highest node where it bends, and the longest path bending at a node is its left depth plus its right depth. I do a post-order DFS where each call returns its depth, one plus the larger child depth, and on the way records left plus right in a running maximum. That visits each node once, O(n) time and O(h) space. The return value is a depth, because a parent can only extend one branch; the diameter is the side result."

Depth and balance checks with the same bottom-up pattern are in tree depth and balance, and validate a binary search tree is another problem where the recursion passes information the answer does not show directly.