Tree diameter: return one branch, combine two locally

Separate the information returned upward from the answer here

A baseline considers every node as a possible highest point of a path and recomputes the heights of both child subtrees. It is correct, but repeated height walks cost O(n²) on a chain. The useful change is postorder traversal: solve children before their parent and retain their height return values.

Define height as the number of nodes on a longest downward path from a node; an empty child has height zero. If left and right heights are L and R, the longest path whose highest node is current has L + R edges: each child height counts its connecting edge to the current node plus its remaining downward edges. The height returned to the parent is 1 + max(L, R), because a path extended by a parent can continue through only one child branch.

Returning L + R + 1 as height is wrong: it would send a branching shape upward and later count a path that visits the branch point twice. Keep a separate maximum diameter, or return (height, best_diameter) with best=max(left_best, right_best, L+R). Every simple tree path has exactly one highest node relative to the root, so checking every node's crossing candidate covers paths entirely inside subtrees as well as those through the root.

The reference simulates recursive entry/return with an explicit expanded flag, stores child heights, and updates one best value. It validates identities on entry so malformed graphs cannot be treated as trees. Each node is processed twice, giving O(n) time. Retained heights and structural identities cost O(n) auxiliary space, plus O(h) frame space; output is O(1). A trusted recursive variant can use O(h) auxiliary space but may exceed Python's recursion budget. Node values and equality never affect topology or distance.

Follow each child height to its parent. Predict the returned height and local diameter before the second return arrives; their units and purposes differ.

Follow-up 1: return endpoint identities as well as distance

Predict the additional return state: each downward height must also name its deepest endpoint. When L + R improves the best, record the two branch endpoints, using current itself for an empty side.

Define a tie rule using stable IDs or traversal order; comparing opaque values cannot enforce identity ordering. A path with k edges has k+1 returned nodes.

Follow-up 2: edges carry nonnegative weights

The edge-count diameter is three, but weighted distance b–a–c–d is ten. Each returned branch becomes edge_weight + child_downward_distance; take the best two branches locally and one upward. Empty downward distance is zero, replacing node-count height. Allowing negative weights requires a new policy on empty paths and endpoint distinctness before using zero as a fallback. A senior candidate derives units and distinguishes returned height from local diameter. A lead candidate defines weight validation, tie semantics, and snapshot ownership; adding cross-links destroys the unique-path proof and requires a graph problem.

Run and check

From the repository root:

cd curriculum/01-code/02-data-structures-algorithms/problems/20-tree-diameter
python -m unittest -v test_solution.py

Reference implementation (download file, source below) · Contract and oracle tests (download file, source below). Read the tests after your attempt. A green reference suite verifies the supplied implementation; it does not demonstrate independent transfer. Reimplement one follow-up with the reference closed and explain which old invariant no longer holds.

from dataclasses import dataclass

@dataclass(eq=False)
class Node:
    value: object
    left: "Node | None" = None
    right: "Node | None" = None

def tree_diameter(root):
    if root is None:
        return 0
    stack = [(root, False)]
    seen, heights = set(), {}
    best = 0
    while stack:
        node, expanded = stack.pop()
        if not expanded:
            if not isinstance(node, Node) or node in seen:
                raise ValueError("expected a tree without malformed or repeated nodes")
            seen.add(node)
            stack.append((node, True))
            if node.right is not None:
                stack.append((node.right, False))
            if node.left is not None:
                stack.append((node.left, False))
            continue
        left = heights.get(node.left, 0)
        right = heights.get(node.right, 0)
        heights[node] = 1 + max(left, right)
        best = max(best, left + right)
    return best

import random
import unittest
from collections import deque
from solution import Node, tree_diameter

class DiameterTests(unittest.TestCase):
    def test_units_and_deep_chain(self):
        self.assertEqual(tree_diameter(None), 0)
        self.assertEqual(tree_diameter(Node(1)), 0)
        root = Node("a", Node("b", Node("d"), Node("e")), Node("c"))
        self.assertEqual(tree_diameter(root), 3)
        chain = None
        for _ in range(5000):
            chain = Node(7, chain)
        self.assertEqual(tree_diameter(chain), 4999)
        # Diameter is entirely below root, so root-height alone is insufficient.
        fork = Node(0, Node(0, Node(0, Node(0))), Node(0, right=Node(0, right=Node(0))))
        self.assertEqual(tree_diameter(Node(0, fork)), 6)

    def test_invalid_structure(self):
        child = Node(1)
        root = Node(0)
        root.left = root
        for invalid in [root, Node(0, child, child), Node(0, right=42), 42]:
            with self.assertRaises(ValueError):
                tree_diameter(invalid)

    def test_all_pairs_distance_oracle(self):
        rng = random.Random(120)
        for _ in range(150):
            nodes = [Node(7) for _ in range(rng.randrange(1, 30))]
            adjacent = {node: [] for node in nodes}
            slots = [(nodes[0], "left"), (nodes[0], "right")]
            for child in nodes[1:]:
                parent, side = slots.pop(rng.randrange(len(slots)))
                setattr(parent, side, child)
                adjacent[parent].append(child)
                adjacent[child].append(parent)
                slots.extend([(child, "left"), (child, "right")])
            expected = 0
            for start in nodes:
                queue, seen = deque([(start, 0)]), {start}
                while queue:
                    node, distance = queue.popleft()
                    expected = max(expected, distance)
                    for neighbor in adjacent[node]:
                        if neighbor not in seen:
                            seen.add(neighbor)
                            queue.append((neighbor, distance + 1))
            self.assertEqual(tree_diameter(nodes[0]), expected)

if __name__ == "__main__":
    unittest.main()