Diameter of Binary Tree
Easy · Trees · Explanation and code example
Iterative DFS
How it works
Recursive DFS is the easiest way to compute diameter, but it uses the call stack. We can simulate the same behavior iteratively using our own stack.
We perform a post-order traversal:
- Visit left subtree
- Visit right subtree
- Then process the current node
For each node, we store in a map:
- its height
- its best diameter
After both children are processed, we can compute:
height = 1 + max(leftHeight, rightHeight)diameter = max(leftHeight + rightHeight, leftDiameter, rightDiameter)
This means every node is processed exactly once.
Steps
- Use a stack to simulate DFS.
- Maintain a map storing
(height, diameter)for each visited node. - For each node:
- First push its children until you reach the bottom (post-order).
- When both children are processed, pop the node:
- Retrieve left and right heights/diameters.
- Compute
height = 1 + max(leftHeight, rightHeight)anddiameter = max(leftHeight + rightHeight, leftDiameter, rightDiameter). - Save these results in the map.
- The final diameter is the second value stored for the root.
Complexity
- Time complexity:
O(n) - Space complexity:
O(n)
Things to watch for
Returning Diameter Instead of Height from DFS
The DFS function must return height (for the parent to use), but update the global diameter. A common mistake is returning the diameter, which breaks the height calculation for ancestors.
Assuming the Longest Path Goes Through the Root
The diameter may pass through any node, not just the root. A solution that only calculates leftHeight + rightHeight at the root will fail on trees where the longest path is entirely within a subtree.
Python code example
from typing import Optional
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, val=0, left=None, right=None):
# self.val = val
# self.left = left
# self.right = right
class Solution:
def diameterOfBinaryTree(self, root: Optional[TreeNode]) -> int:
stack = [root]
mp = {None: (0, 0)}
while stack:
node = stack[-1]
if node.left and node.left not in mp:
stack.append(node.left)
elif node.right and node.right not in mp:
stack.append(node.right)
else:
node = stack.pop()
leftHeight, leftDiameter = mp[node.left]
rightHeight, rightDiameter = mp[node.right]
mp[node] = (1 + max(leftHeight, rightHeight),
max(leftHeight + rightHeight, leftDiameter, rightDiameter))
return mp[root][1]
C++ code example
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode() : val(0), left(nullptr), right(nullptr) {}
* TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
* TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
* };
*/
class Solution {
public:
int diameterOfBinaryTree(TreeNode* root) {
unordered_map<TreeNode*, pair<int, int>> mp;
mp[nullptr] = {0, 0};
stack<TreeNode*> stack;
stack.push(root);
while (!stack.empty()) {
TreeNode* node = stack.top();
if (node->left && mp.find(node->left) == mp.end()) {
stack.push(node->left);
} else if (node->right && mp.find(node->right) == mp.end()) {
stack.push(node->right);
} else {
node = stack.top();
stack.pop();
auto[leftHeight, leftDiameter] = mp[node->left];
auto[rightHeight, rightDiameter] = mp[node->right];
int height = 1 + std::max(leftHeight, rightHeight);
int diameter = max(leftHeight + rightHeight,
max(leftDiameter, rightDiameter));
mp[node] = {height, diameter};
}
}
return mp[root].second;
}
};
Source and license
Adapted from NeetCode’s source · MIT license. Node types are supplied by LeetCode where needed.