Diameter of Nary tree using Dynamic Programming
Get this book > Problems on Array: For Interviews and Competitive Programming
In this problem, we are given input as the reference to the root of an Nary tree.
We need to calculate the diameter of the tree that is the longest path between any two nodes using Dynamic Programming.
A tree is an extensively used data structure in the world of programming. Every node in a tree can have more further subdivisions. The bottommost nodes of a tree that have no subdivisions, are called the leaf nodes.
We know that in a binary tree, each node has no more than 2 children. Similarly, in an Nary tree, each node has no more than N children. N can be any integer here.
DIAMETER OF A TREE: It is defined as the number of nodes on the longest path between 2 nodes.This path may or may not pass through the root of the tree.This path includes two leaf nodes. It is also known as the width of a tree.
Example 1: diameter passing through the root node:
Example 2: diameter doesn't pass through the root node:
To solve this, we need to proceed by breaking it down into subproblems that follow a similar structure and storing the values for future use. This is the approach used for Dynamic Programming.
We will be using DFS i.e Depth First Search algorithm to solve this problem. This algorithm starts at the root node of the tree and explores as far as possible along each branch before backtracking. Therefore, we start with the leaf nodes and make our way to the top.
There are 4 main variables used in this implementation:

FirstMax: This variable stores the height of the largest subtree of any node.

SecondMax: This variable stores the height of the second largest subtree of any node.
In this example, for node B, the largest subtree is D>E>H. Therefore, FirstMax = 3
The second largest subtree is C>F (or C>A). Therefore, SecondMax = 2.
For any node, by default, FirstMax and SecondMax are 1. When we discover the subtrees of the nodes through DFS, FirstMax and SecondMax are changed accordingly.
 dp1[node] : This calculates the no. of nodes that will be in the diameter of the tree, if we consider, that for any node, the diameter doesn't pass through 2 of its subtrees.
In the given example,
the diameter passes through B>A>C>D. For node C, the diameter doesn't pass through 2 of its subtrees.It only passes through one subtree i.e. D. Hence, till node C, 2 nodes(C and D) will be in the diameter of the tree. dp1[C] = 2
dp1[node] = 1 + FirstMax
 dp2[node]: This calculates the no. of nodes that will be in diameter of the tree, if for any node, the diameter passes through two of its subtrees.
In this example,
the diameter passes through E>D>B>C>F.
For node B, the diameter passes through 2 of its subtrees. Therefore, 5 nodes are in the diameter.
dp2[B] = 5
dp2[node] = 1 + FirstMax + SecondMax
For a leaf node, we consider dp1[node]=1 and dp2[node] = 0.
There is also an extra variable max_diam, that keeps a track of the maximum diameter till that point, so as to attain the final diameter of the tree.
Let us walk through this problem with a basic example:
We traverse through the whole tree using DFS by starting at the leaf nodes and making our way to the top. The Firstmax, SecondMax, dp1[node] and dp2[node] are calculated along the way.
Here max_diam will be 6, which will be returned to the main function.
Therefore, the diameter of the given tree is 6.
Below is the implementation of the above approach in C++:
#include <bits/stdc++.h>
using namespace std;
int diameter = 1;
int max_diam = 0;
// Function to find the diameter of the tree
// using Dynamic Programming
int dfs(int node, int parent, int dp1[], int dp2[], list<int>* adj)
{
// Store the first maximum and secondmax
int firstmax = 1;
int secondmax = 1;
// cout<<"\nmax_diam is "<<max_diam;
// Traverse for all children of node
for (auto i = adj[node].begin(); i != adj[node].end(); ++i) {
if (*i == parent)
continue;
// cout<<"\nbefore dfs node is "<<node;
// Call DFS function again
dfs(*i, node, dp1, dp2, adj);
// cout<<"\nafter dfs node is "<<node;
// Find first max
if (firstmax == 1) {
firstmax = dp1[*i];
}
else if (dp1[*i] >= firstmax) // Secondmaximum
{
secondmax = firstmax;
firstmax = dp1[*i];
}
else if (dp1[*i] > secondmax) // Find secondmaximum
{
secondmax = dp1[*i];
}
}
// Base case for every node
dp1[node] = 1;
if (firstmax != 1) // Add
dp1[node] += firstmax;
// Find dp[2]
if (secondmax != 1)
dp2[node] = 1 + firstmax + secondmax;
// Return maximum of both
if (max(dp1[node], dp2[node])>max_diam)
max_diam = max(dp1[node], dp2[node]);
return max_diam;
}
// Driver Code
int main()
{
int n = 8;
/* Constructed tree is
1
/ \ \
2 3 4
/ \ \
5 6 7
/
8
*/
list<int>* adj = new list<int>[n + 1];
/*create undirected edges */
adj[1].push_back(2);
adj[2].push_back(1);
adj[1].push_back(3);
adj[3].push_back(1);
adj[1].push_back(4);
adj[4].push_back(1);
adj[2].push_back(5);
adj[5].push_back(2);
adj[2].push_back(6);
adj[6].push_back(2);
adj[3].push_back(7);
adj[7].push_back(3);
adj[6].push_back(8);
adj[8].push_back(6);
int dp1[n + 1], dp2[n + 1];
memset(dp1, 0, sizeof dp1);
memset(dp2, 0, sizeof dp2);
// Find diameter by calling function
cout << "Diameter of the given tree is "
<< dfs(1, 1, dp1, dp2, adj) << endl;
return 0;
}
Output:
Diameter of the given tree is 6
Time complexity: O(n)
Conclusion:
We have studied one of the many approaches for calculating the diameter of an Nary tree using Dynamic Programming.