Problem overview
The method diameter receives the root of a binary tree and returns the length of its diameter: the longest path between any two nodes. The length counts edges, not nodes, and the path may or may not pass through the root.
For the tree with root 1, left child 2 (children 4 and 5) and right child 3, the answer is 3. The path 4, 2, 1, 3 has four nodes and three edges, and so does 5, 2, 1, 3. For the tree with root 1 and a single child 2, the answer is 1.
Approach: every subtree reports two numbers
Look at one node and the longest path whose highest node is that one. 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, where each side's height includes the edge from the node to its child. The diameter of the whole tree is the best such path over all nodes.
A direct implementation computes the height of the two subtrees at every node, and that repeats work. Instead, one recursive call per node can return both numbers its parent needs: the height of the subtree and the best diameter found inside it. The solution packs them into an Array[Int] of two elements, with the height at index 0 and the diameter at index 1.
At each node, with leftHeight and rightHeight counting edges down to the deepest node on each side:
- the height is the larger of the two;
- the diameter is the best of the left subtree's diameter, the right subtree's diameter and the bent path through this node,
leftHeight + rightHeight.
For the first example, the leaves 4, 5 and 3 return [0, 0]. Node 2 has two children, so both heights become 1, its height is 1 and its diameter is 2: [1, 2]. At the root, the left side has height 1 + 1 = 2 and the right side has 0 + 1 = 1. The root's diameter is the best of 2, 0 and 2 + 1 = 3, which is 3.
Step by step
- Use the helper in
diameter. Calldfs(root)and take element 1 of the result, the diameter. Arrays are indexed with parentheses, so this isdfs(root)(1). - Declare
dfswith aTreeNodeparameter and the result typeArray[Int]. - Base case. For a
nullnode the subtree is empty, so returnArray(0, 0). - Recurse left. Keep the array in
left, then readleftHeightas avar, because it will change, andleftDiameteras aval. - Recurse right in the same way for
rightHeightandrightDiameter. - Add the edge to the left child. Add 1 to
leftHeightwhenroot.leftis not null, using anifexpression in parentheses. The height the child reports counts edges inside the child's subtree, and the edge from this node to the child is one more. - Add the edge to the right child the same way. With no right child nothing is added, since no edge goes that way.
- Compute the height as
math.max(leftHeight, rightHeight). - Compute the diameter with two nested
math.maxcalls overleftDiameter,rightDiameterandleftHeight + rightHeight, becausemath.maxtakes two arguments. - Return
Array(height, diameter), in that order, as the last expression.
The solution in Scala
class Solution {
def dfs(root: TreeNode): Array[Int] = {
if (root == null) {
return Array(0, 0)
}
val left = dfs(root.left)
var leftHeight = left(0)
val leftDiameter = left(1)
val right = dfs(root.right)
var rightHeight = right(0)
val rightDiameter = right(1)
leftHeight += (if (root.left != null) 1 else 0)
rightHeight += (if (root.right != null) 1 else 0)
val height = math.max(leftHeight, rightHeight)
val diameter = math.max(math.max(leftDiameter, rightDiameter), leftHeight + rightHeight)
Array(height, diameter)
}
def diameter(root: TreeNode): Int = {
dfs(root)(1)
}
}
Scala details:
- Array indexing uses parentheses:
left(0), notleft[0].dfs(root)(1)calls the method first and then indexes the array it returns. ifis an expression. Scala has no ternary operator, and(if (cond) 1 else 0)is its form. The parentheses make it one operand of+=.varandval. Only the two heights are reassigned, so only they arevar. The rest areval.- The names. Inside
dfs, the localval diameteris a different thing from the methoddiameterof the class. The local name wins insidedfs, and the method is unaffected. math.maxcomes from thescala.mathpackage, which is available without an import. It takes exactly two arguments.- Return. The early
return Array(0, 0)is allowed in a method with a declared result type. The final expressionArray(height, diameter)is returned without the keyword.
Complexity
Time: O(n). Each node is visited once and does a constant amount of work, plus one constant-time call for each missing child.
Space: O(h), where h is the height of the tree. The recursion goes as deep as the tree, and each call creates a small two-element array, which becomes garbage after its parent has read it. For a balanced tree that is about log n, but for a chain of nodes it is n, and a very long chain can exhaust the JVM's call stack with a StackOverflowError.
Edge cases and common mistakes
- A single node. The diameter is 0, since no edge exists. The root has two empty children, no edge is added, and
0 + 0is returned. - Counting nodes instead of edges. A path of four nodes has length 3. The solution gives a leaf height 0 and adds an edge only when a child exists. Counting nodes would be one too large.
- Forgetting the +1 for the edge to the child. Then a node with two leaf children would get diameter 0 instead of 2.
- The longest path need not pass through the root. When one subtree is much bigger than the other, the best path can lie entirely inside it. That is why the diameter takes the maximum with
leftDiameterandrightDiameter. In the last test below, the best path has 4 edges while the longest path through the root has only 3. - A chain. A chain of four nodes, each the left child of the previous one, has diameter 3, the height of the chain.
- Mixing up the array order. The caller reads the height at index 0 and the diameter at index 1, so the helper must build the array in that order.
Testing
- Root 1, left child 2 (children 4, 5), right child 3: 3.
- Root 1 with a single child 2: 1.
- A single node: 0.
- A chain of four nodes: 3.
- Root 1 whose only child is 2, where 2 has the children 3 and 4, 3 has a child 5 and 4 has a child 6: 4, along the path 5, 3, 2, 4, 6, which avoids the root.
