Diameter of Binary Tree

You'll compute every subtree's height once, children first, and check the longest path bending at each node. Read Balanced Binary Tree first; it uses the same traversal. Related chapter: trees.

1. The problem

You get the root of a binary tree. Return the most edges on a path between any two nodes. This is called the diameter. The path doesn't have to pass through the root. The values don't matter. Don't change the tree. Our function also returns 0 for an empty tree. This is LeetCode 543.

"Edges" is the word that matters: a path through three nodes has two edges.

2. A correct baseline

Turn the tree into a graph where each node lists its neighbors. Then walk the whole graph from every node, and keep the longest distance found.

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

def example():
    return TreeNode(1, TreeNode(2, TreeNode(4), TreeNode(5)), TreeNode(3))

def baseline(root):
    if root is None:
        return 0
    graph = {}
    pending = [root]
    while pending:
        node = pending.pop()
        graph.setdefault(id(node), [])
        for child in (node.left, node.right):
            if child is not None:
                graph.setdefault(id(child), []).append(id(node))
                graph[id(node)].append(id(child))
                pending.append(child)
    best = 0
    for start in graph:
        pending = [(start, None, 0)]
        while pending:
            node, parent, distance = pending.pop()
            best = max(best, distance)
            for neighbor in graph[node]:
                if neighbor != parent:
                    pending.append((neighbor, node, distance + 1))
    return best

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

Each of the n walks visits all n nodes, so this takes O(n²) time. On 10,000 nodes, that's 100 million steps. It uses O(n) extra space.

The bottleneck. Each walk explores the same branches again, just from a different starting node.

3. The observation

In the example, the longest path is 4 → 2 → 1 → 3. It goes up to node 1 and bends back down. Every path has one highest node where it bends like this.

The longest path that bends at a node goes down its left side as far as possible, and down its right side as far as possible. In the example, node 1's left subtree has height 2 and its right has height 1. Heights count nodes, and each node on a side adds one edge to the path. So the path has 2 + 1 = 3 edges.

So compute every subtree's height once, children before parents, as in Balanced Binary Tree. At each node, check left + right, and keep the largest.

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 diameterOfBinaryTree(self, root) -> int:
        """Return the longest path's edge count; empty trees return zero."""
        heights = {id(None): 0}
        pending = [(root, False)]
        best = 0
        while pending:
            node, expanded = pending.pop()
            if node is None:
                continue
            if not expanded:
                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)]
                best = max(best, left + right)  # Path bending at this node.
                heights[id(node)] = 1 + max(left, right)
        return best

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

Trace the example. Each row is a node whose children are done:

5. Why it works

What stays true. Children finish before their parent, so every saved height is correct.

best is never too big. Joining the deepest path on each side of a node makes a real path with left + right edges.

best is never too small. The longest path bends at some highest node. Neither of its two sides can be longer than that node's child heights. So the check at that node is at least as long as the longest path. Together, best is exactly the diameter.

6. Cost, edge cases, and failure

Cost. Each node is handled twice: O(n) time. The saved heights take O(n) space, and the stack O(h). Using your own stack avoids Python's recursion limit.

Edge cases. One node has diameter 0. Two nodes have diameter 1, because we count edges, not nodes.

Where it fails. Checking only the path through the root can miss the answer:

fork = TreeNode(2, TreeNode(3, TreeNode(4)), TreeNode(5, None, TreeNode(6)))
root = TreeNode(1, fork)
assert solve(root) == 4  # 4 → 3 → 2 → 5 → 6, entirely below the root.
assert solve(TreeNode(1, TreeNode(2))) == 1  # Count edges, not nodes.

import random
rng = random.Random(543)
def random_tree(depth):
    if depth == 0 or rng.randrange(3) == 0:
        return None
    return TreeNode(rng.randrange(3), 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):
    root = random_tree(6)
    before = snapshot(root)
    assert solve(root) == baseline(root)
    assert all(node.val == val and node.left is left and node.right is right
               for node, val, left, right in before)
chain = None
for _ in range(10000):
    chain = TreeNode(1, chain)
assert solve(chain) == 9999