Problem overview
The diameter of a binary tree is the length of the longest path between any two of its nodes, measured in edges. The path may or may not pass through the root. The function diameter returns it for a tree stored in arrays: left[i] and right[i] are the indices of the left and right children of node i, root is the index of the root, and the index 0 means that there is no node. Real nodes therefore have indices from 1 upwards.
Take the tree with root 1, children 2 and 3, and 2's children 4 and 5. In the arrays it looks like this:
left = [0, 2, 4, 0, 0, 0]
right = [0, 3, 5, 0, 0, 0]
root = 1The longest path is 4, 2, 1, 3 (or 5, 2, 1, 3), which has three edges, so the answer is 3. A tree with just a root and one child has diameter 1, and a lone root has diameter 0.
Approach: every node is the bend of one candidate path
The highest node on any path is the point where it bends: from there it goes down into the left subtree, the right subtree, or both. Suppose the longest downward path in the left subtree has lh nodes and the longest one in the right subtree has rh nodes (these are the subtrees' heights, counted in nodes). Reaching the far end of the left path takes lh edges from the current node, and the right path takes rh edges. So the longest path that bends at this node has lh + rh edges. Because a path can bend at any node, the diameter is the largest value of lh + rh over all nodes.
A slow way computes the height of both subtrees from scratch at every node, which costs O(n²) for a long chain. The fast way computes everything in a single post-order traversal. Each call returns two numbers for its subtree, its height and the best diameter found inside it, and the parent combines them:
- height = the taller child's height + 1;
- diameter = the best of the left subtree's diameter, the right subtree's diameter, and the bent path
lh + rhthrough this node.
An empty subtree has height 0 and diameter 0. With this convention a leaf has height 1, and heights count nodes, which is why lh + rh is already the number of edges. For example, a node with a leaf on each side has lh = 1 and rh = 1, so the path leaf, node, leaf has two edges.
Step by step
- Delegate to
dfs. Indiameter, callSolution::dfs(left, right, root)and take field.1of the returned tuple, which is the diameter. - Declare
dfswith the two child slices and the currentnodeindex. It returns a tuple of twoi32values: the height and the diameter. - Base case. When
nodeis 0, return(0, 0). - Left subtree: call
dfsonleft[node]and destructure the result intolhandld. - Right subtree: do the same on
right[node], givingrhandrd. - Height:
lh.max(rh) + 1. - Diameter:
ld.max(rd).max(lh + rh). - Return the tuple
(height, diameter)as the last expression. The order matters because the caller reads the diameter from.1.
The solution in Rust
impl Solution {
fn dfs(left: &[usize], right: &[usize], node: usize) -> (i32, i32) {
if node == 0 {
return (0, 0);
}
let (lh, ld) = Solution::dfs(left, right, left[node]);
let (rh, rd) = Solution::dfs(left, right, right[node]);
let height = lh.max(rh) + 1;
let diameter = ld.max(rd).max(lh + rh);
(height, diameter)
}
fn diameter(left: &[usize], right: &[usize], root: usize) -> i32 {
Solution::dfs(left, right, root).1
}
}
Rust details that matter here:
- Slices as parameters.
left: &[usize]andright: &[usize]borrow the arrays without copying them, so every recursive call can share the same data. Both a vector reference and an array reference can be passed to them. - Indices are
usize. Rust indexes slices withusize, so the child indices andnodeuse that type. A zero is a valid value there, and index 0 of each slice is unused padding that stands for "no node". - Returning a tuple.
(i32, i32)lets one call report two facts.let (lh, ld) = ...destructures it, and.1reads the second field by position. - Chained
max.ld.max(rd).max(lh + rh)is the largest of three values.maxis a method ofi32that takes the other value. - Calling through
Solution::. The functions take noself, so they are associated functions, called with the type's name as a prefix. - Tail expressions.
(height, diameter)andSolution::dfs(left, right, root).1have no semicolon, so they are the return values. The earlyreturn (0, 0);in the base case does needreturnand its semicolon.
Complexity
Time: O(n). Each node index is visited once, and each visit does a few comparisons. Calls on index 0 add at most n + 1 more.
Space: O(h). Beyond the input arrays, the only memory is the call stack, as deep as the tree is tall.
Edge cases and common mistakes
- Single node. With
left = [0, 0],right = [0, 0]androot = 1, both subtrees return(0, 0), so the diameter is0 + 0 = 0. - Root and one child. With
left = [0, 2, 0]andright = [0, 0, 0], the answer is 1. - A path that doesn't touch the root. Take the arrays
left = [0, 2, 3, 5, 6, 0, 0]andright = [0, 0, 4, 0, 0, 0, 0]. Node 1 has only the left child 2, node 2 has children 3 and 4, and 3 and 4 each have one left child, 5 and 6. The longest path is 5, 3, 2, 4, 6 with four edges, and it never reaches the root. Taking onlylh + rhat the root would give 3, which is too small. That is why the diameters of the subtrees,ldandrd, take part in the maximum. - Off-by-one with heights. The base case returns 0, which makes a leaf's height 1 and keeps
lh + rhequal to the edge count. If empty subtrees returned -1, thenlh + rhwould have to becomelh + rh + 2. Mixing the two conventions gives answers that are off by one or two. - Using the height instead of the diameter. Returning the height from
diameteris correct only if the longest path happens to go through the root. - Chains. For a tree shaped like a line, such as
left = [0, 2, 3, 0]andright = [0, 0, 0, 0]withroot = 1, the diameter is the number of edges in the chain, here 2.
Testing
left = [0, 2, 4, 0, 0, 0],right = [0, 3, 5, 0, 0, 0],root = 1gives 3.left = [0, 2, 0],right = [0, 0, 0],root = 1gives 1.left = [0, 0],right = [0, 0],root = 1gives 0.left = [0, 2, 3, 5, 6, 0, 0],right = [0, 0, 4, 0, 0, 0, 0],root = 1gives 4.left = [0, 2, 3, 0],right = [0, 0, 0, 0],root = 1gives 2.
