Site icon Study Algorithms

Diameter of a Binary Tree

You are provided with a binary tree, return the length of the diameter. The diameter of the binary tree is defined as the length of the longest path between any two nodes in the tree.

Diameter of a binary tree 7 through root

In the above example the diameter of the binary tree is 7.

Problem Statement:

We can easily find out the distance between two nodes in a binary tree. In the given example, the distance between node 5 and 8 is 4. The distance between node 2 and 6 is 3, the distance between node 8 and 0 is 2 and so on. To put the problem in simple words, the diameter of a binary tree is the maximum of all such lengths you can find. If you try to calculate all the distances between every possible node, you would see that longest path between 2 nodes is 7. (formed by nodes 6 and 9).

Fig: Different path lengths and diameter
Brute Force Method:

Just as you look at the problem, one solution seems very obvious.

But this method would only work efficiently for small trees having 5-6 nodes. Once the tree begins to grow in size, you realize that we are wasting a lot of computation in calculating unnecessary paths making the time complexity to be O(n^2).

Recursive Method:

A very popular way to solve this problem is using recursion. The way recursion works is that every left and right child of a node is a tree in itself. So the maximum possible path can be found in the left sub-tree or in the right sub-tree. It may also be possible that the path includes the root node.

So, if the longest path will include the root node, then the longest path must be the depth(root->right) + depth (root->left). If not, then we would find the diameter either in the left sub-tree or in the right sub-tree. The algorithm for a recursive method looks like:

Code:
int max = 0;
private int diameter(TreeNode root) {
    max = 1;
    depth(root);
    return max - 1;
}

private int depth(TreeNode node)
{
    if(node == null)
    {
        return 0;
    }

    int leftDepth = depth(root.left);
    int rightDepth = depth(root.right);

    // update the answer, because diameter of a  
    // tree is nothing but maximum value of  
    // (left_depth + right_depth + 1) for each node  
    max = Math.max(max, leftDepth + rightDepth + 1);
    return 1 + Math.max(leftDepth, rightDepth);        
}
Code language: Java (java)

This method works perfectly and has a time complexity of O(n). But it occupies stack space in memory and can be harder to debug because of the recursion involved.

Iterative Method:

To this point, we are clear about one thing. To find the diameter of a tree, we need to process the left and right sub-tree first. This implies that using a post-order traversal may be helpful here. It makes sure that we traverse the left and right sub-tree before moving to the root.

Here is an image showing maximum depth of each of the nodes in the tree. The diameter of the tree is the maximum value of (left_depth + right_depth) among each of the nodes.

Fig: Showing maximum depth at each node.
Diameter is the maximum of (left_depth + right_depth) = 7

We use stacks to perform a post order traversal of the binary tree. A hash-map can be used to maintain the maximum depth at each of the nodes in the tree.

Code:
  public int diameterOfBinaryTree(TreeNode root) {

    Map<TreeNode, Integer> map = new HashMap<>();
    Stack<TreeNode> stack = new Stack<>();
    int diameter = 0;

    if (root != null)
      stack.push(root);

    while (!stack.isEmpty()) {
      TreeNode node = stack.peek();

      // Fill up stack to perform post-order traversal
      if (node.left != null && !map.containsKey(node.left)) {
        stack.push(node.left);
      } else if (node.right != null && !map.containsKey(node.right)) {
        stack.push(node.right);
      } else {

        // Process the root, once left and right sub-tree have been processed
        stack.pop();
        int leftDepth = map.getOrDefault(node.left, 0);
        int rightDepth = map.getOrDefault(node.right, 0);

        // Put the max depth at a node in the map
        map.put(node, 1 + Math.max(leftDepth, rightDepth));

        // Update the max diameter found so far
        diameter = Math.max(diameter, leftDepth + rightDepth);
      }
    }
    return diameter;
  }
Code language: Java (java)

Time Complexity: O(n)
Space Complexity: O(n)

You can find a working solution here.
The code and test cases are also available on Github.
A similar problem can be found on Leetcode.

Video Explanation:
Exit mobile version