Problem overview
The function diameter receives the root of a binary tree, a *TreeNode, and returns the length of its diameter: the longest path between any two nodes, counted in edges. The path may or may not pass through the root. The tree has between 1 and 10,000 nodes.
For the tree [1,2,3,4,5], node 1 has children 2 and 3, and node 2 has children 4 and 5. The path 4, 2, 1, 3 uses three edges, and so does 5, 2, 1, 3. No path is longer, so the answer is 3. For [1,2] the only path is the single edge between the two nodes, so the answer is 1.
Approach: combine heights on the way back up
Think about the longest path that bends at one particular node. It goes down into the left subtree as far as it can, and down into the right subtree as far as it can. Its length is the height of the left side plus the height of the right side, so every node gives one candidate answer. The diameter is the best candidate over all nodes.
The straightforward approach computes the height of both subtrees at every node, which repeats work and takes O(n²) time on a tree shaped like a chain. The better idea is to let each recursive call report two numbers: the height of its subtree and the best diameter found inside it. A node then needs only the numbers its two children returned:
- its height is the taller child's height plus one;
- its diameter is the largest of the left diameter, the right diameter, and the path through the node itself.
Heights here count nodes, so an empty subtree has height 0 and a leaf has height 1. With that convention, leftHeight + rightHeight equals the number of edges on the path that bends at the node. The longest chain below the bend node holds leftHeight nodes on the left and rightHeight on the right, so the whole path has leftHeight + rightHeight + 1 nodes, and a path with that many nodes has leftHeight + rightHeight edges.
For [1,2,3,4,5]: node 4 and node 5 are leaves with height 1 and diameter 0. Node 2 has height 2 and diameter 1 + 1 = 2. Node 3 has height 1. At the root, the path through it is 2 + 1 = 3, which beats the left diameter 2 and the right diameter 0, so the answer is 3.
Step by step
- In
diameter, calldiameterDfs(root). It returns the height and the diameter. The height isn't needed here, so assign it to the blank identifier_, which Go accepts for values you must receive but don't use. Store the diameter inresultand return it. - Declare
diameterDfswith a*TreeNodeparameter and twointresults, written(int, int). - Base case. For a
nilnode, return0, 0: no nodes, no path. - Recurse left.
leftHeight, leftDiameter := diameterDfs(root.Left)declares both variables at once. - Recurse right in the same way to get
rightHeightandrightDiameter. - Compute
heightas the larger child height plus 1, using the built-inmax. - Compute
resultas the largest of three candidates: the left diameter, the right diameter, andleftHeight+rightHeight. - Return
height, result, in the order the caller unpacks them.
The solution in Go
func diameter(root *TreeNode) int {
_, result := diameterDfs(root)
return result
}
func diameterDfs(root *TreeNode) (int, int) {
if root == nil {
return 0, 0
}
leftHeight, leftDiameter := diameterDfs(root.Left)
rightHeight, rightDiameter := diameterDfs(root.Right)
height := max(leftHeight, rightHeight) + 1
result := max(max(leftDiameter, rightDiameter), leftHeight+rightHeight)
return height, result
}
The Go details worth a closer look:
- Multiple return values. A Go function can return several values, and the caller must receive all of them.
_throws one away. Without it, an unused variable such asheightwould be a compile error. :=with two names.leftHeight, leftDiameter := ...declares both variables from the two results.- Built-in
max. Since Go 1.21,maxis built in and takes two or more numbers. Older code wrote its own helper. Here it is called with two arguments each time, somax(max(a, b), c)is the largest of three.max(a, b, c)would also work. - Return order. The results are positional:
diameterDfsreturns the height first, so_, result :=keeps the second value.
Complexity
Time: O(n). Each node is visited once and does a constant amount of work with the numbers its children return.
Space: O(h) for the recursion, where h is the height of the tree: about log n when balanced, and up to n when every node has one child. With at most 10,000 nodes, even the chain case is a modest depth for Go, whose stacks grow as needed.
Edge cases and common mistakes
- A single node.
[1]gives 0: both heights are 0, so the path through the root has no edges. - The longest path avoids the root. If one subtree is a large bushy tree and the other side is a single leaf, the best path is entirely inside the bushy subtree. Returning only
leftHeight + rightHeightfor the root would miss it, which is why the diameters of the children are candidates too. - A chain. For
[1,2,null,3,null,4], every node has one child. The answer is 3, the length of the whole chain. - Counting nodes instead of edges. The answer is the number of edges. Adding 1 to
leftHeight + rightHeightwould report 4 for[1,2,3,4,5]. - Mixing up the return order. Swapping
heightandresultin the return statement compiles, because both areint, but gives wrong answers. The compiler can't catch it.
Testing
[1,2,3,4,5]gives3.[1,2]gives1.[1]gives0.[1,2,null,3,null,4]gives3.