Problem overview
The method diameter receives the root of a binary tree as a TreeNode? and returns the tree's diameter: the number of edges on the longest path between any two nodes. The path doesn't have to go through the root, and it can go up and then down but never visits a node twice.
For the tree [1, 2, 3, 4, 5], the root 1 has children 2 and 3, and 2 has children 4 and 5. The path 4, 2, 1, 3 has 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: height and diameter from the same pass
Think about where a longest path bends. Every path has a highest node, the one closest to the root, and the path goes down from it into its left subtree and into its right subtree. Its length is the longest downward path on the left plus the longest downward path on the right, joined through that node.
So for each node, the best path whose highest node is that node is leftHeight + rightHeight, where each height counts edges and includes the edge from the node to that child. The diameter of the tree is the largest such value over all nodes. Equivalently, for each subtree, the diameter is the best of three things: the diameter entirely inside the left subtree, the diameter entirely inside the right subtree, or the path that bends at this node.
Computing the heights separately for every node would revisit subtrees again and again, which costs O(n²) for a long chain. The trick is to let one recursive function report both numbers for a subtree, a height and a diameter, as a pair. A node then builds its pair from its children's pairs, and each node is visited once.
On [1, 2, 3, 4, 5], the leaves 4, 5 and 3 each report [0, 0]. Node 2 sees two children, so its left and right heights are both 0 + 1 = 1. Its height is 1, and its diameter is the best of 0, 0 and 1 + 1 = 2, which is 2. The root sees node 2 with height 1, plus the edge to it: 2. It sees node 3 with height 0 plus its edge: 1. The path through the root is 2 + 1 = 3, which beats the diameter 2 inside the left subtree. The root's pair is [2, 3], and diameter returns index 1.
Step by step
- Call the helper from
diameter:dfs(root)[1]. The helper returns[height, diameter], and index 1 is the diameter. - Declare
dfsabove it. It takes aTreeNode?and returns an[Int]with two entries. - Base case. An empty subtree has height 0 and diameter 0, so return
[0, 0]. - Recurse left and copy the two entries into
leftHeightandleftDiameter. They are declared together withvarbecause the height changes next. - Recurse right and copy the entries into
rightHeightandrightDiameter. - Add the edge to the left child if there is one. The child's height counts edges below the child, so a path through the node needs one more edge to reach it. Use a ternary: add
1whenroot!.leftis notnil, and0otherwise. - Do the same for the right side.
- Compute this node's height as
max(leftHeight, rightHeight). These values already include the edges to the children. - Compute the diameter as the maximum of
leftDiameter,rightDiameterandleftHeight + rightHeight. Swift'smaxtakes two arguments, so the call is nested. - Return
[height, diameter]. The parent reads the height at index 0 and the diameter at index 1, so the order matters.
The solution in Swift
class Solution {
func dfs(_ root: TreeNode?) -> [Int] {
if root == nil {
return [0, 0]
}
let left = dfs(root!.left)
var leftHeight = left[0], leftDiameter = left[1]
let right = dfs(root!.right)
var rightHeight = right[0], rightDiameter = right[1]
leftHeight += root!.left != nil ? 1 : 0
rightHeight += root!.right != nil ? 1 : 0
let height = max(leftHeight, rightHeight)
let diameter = max(max(leftDiameter, rightDiameter), leftHeight + rightHeight)
return [height, diameter]
}
func diameter(_ root: TreeNode?) -> Int {
return dfs(root)[1]
}
}
Swift details worth a closer look:
- Returning two values in an array. The lesson returns
[Int]and reads it by index. A tuple such as(Int, Int)would also work and could be unpacked by name, but the array keeps this version short. The index meaning (0 for height, 1 for diameter) has to be remembered by the programmer, since the compiler can't check it. - Several declarations in one
var.var leftHeight = left[0], leftDiameter = left[1]declares two variables in one statement. Only the two heights are changed later, soletwould be enough for the diameters, but declaring them withvarworks too. - The ternary operator.
root!.left != nil ? 1 : 0turns a condition into the number 1 or 0. Swift does not convertBooltoIntautomatically, soleftHeight += (condition)would not compile. Precedence helps here: the ternary operator binds more loosely than!=, so no parentheses are needed. maxwith two arguments.max(max(a, b), c)picks the largest of three. Swift's globalmaxalso accepts more arguments in one call, but the lesson nests it, which works the same way.- Optional handling.
root!is safe after thenilcheck. Comparingroot!.leftwithnilasks whether a left child exists without unwrapping it.
Complexity
Let n be the number of nodes and h the height of the tree.
Time: O(n). Every node is visited once, and every visit does a constant amount of work: two small arrays are read, and a few additions and comparisons are made. Each missing child costs one call that returns at once.
Space: O(h). The recursion goes as deep as the tree is tall, which is about log n for a balanced tree and n for a chain. The small arrays that calls return are short-lived.
Edge cases and common mistakes
- A single node. The diameter is 0: there is no edge at all. Both children report
[0, 0], no edge is added, and the result is[0, 0]. - An empty tree. The constraints promise at least one node, but
nilwould also work and return 0. - Counting nodes instead of edges. The path 4, 2, 1, 3 has four nodes and three edges. The answer is 3.
- A diameter that avoids the root. A big subtree hanging off one side can hold the longest path. That is why the diameter comes from a maximum over the children's diameters as well as the bend at the current node.
- Adding the edge unconditionally. If the
+ 1were added even when a child is missing, a leaf would get height 1 and every length would be too big. The ternary prevents that. - Using height alone.
leftHeight + rightHeightat the root only covers paths through the root. Without the stored child diameters, the method would miss paths that stay inside one subtree.
Testing
[1, 2, 3, 4, 5]gives3.[1, 2]gives1.[1]gives0.[1, 2, null, 3, 4, null, null, 5, null, null, 6]gives4, a path that bends below the root.
