DP on Trees | Set3 ( Diameter of Nary Tree ) (original) (raw)
DP on Trees | Set-3 ( Diameter of N-ary Tree )
Last Updated : 11 Jul, 2025
Given an **N-ary tree, the task is to calculate the **longest path between any two nodes (also known as the **diameter of the tree).
**Examples:
**Example1: The diameter of the N-ary tree in the below example is 9, which represents the number of edges along the path marked with green nodes.
**Example2: The diameter of the N-ary tree in the below example is 7, which represents the number of edges along the path marked with green nodes.
**Approach:
The idea is to use Dynamic Programming with two key states: **dp1[node] and **dp2[node]. The value **dp1[node] represents the **longest path from the **current node to **any leaf in its subtree, while **dp2[node] represents the longest path passing through the **current node and two of its subtrees. By combining these two states across all nodes, we compute the **maximum diameter. A Depth First Search (DFS) traversal is used to calculate these states efficiently.
Step-by-step approach:
- Use **dp1[node] (longest path to a leaf in subtree) and **dp2[node] (longest path passing through the node and its two subtrees).
- For leaf nodes, set dp1[node] = 0 and **dp2[node] = 0.
- Begin a DFS traversal from an arbitrary node, tracking the **parent to avoid revisits.
- For each node, find the two largest **dp1[child] values using firstMaxand secondMax.
- Compute **dp1[node] = 1 + firstMax and dp2[node] = firstMax + secondMax.
- Update the **diameter as the maximum of **dp2[node] and the current diameter. C++ `
// C++ program to find the diameter // of n-ary tree using dp. #include <bits/stdc++.h> using namespace std;
class Node { public: int data; vector<Node*> children;
Node(int x) {
data = x;
}};
// Function to perform DFS and calculate dp1 and dp2 int dfs(Node* node, unordered_map<Node*, int>& dp1, unordered_map<Node*, int>& dp2, int& diameter) { if (!node) return 0;
int firstMax = -1, secondMax = -1;
// Traverse children of the current node
for (Node* child : node->children) {
// Recursive call to process child
dfs(child, dp1, dp2, diameter);
// Find first and second maximum dp1 values
if (dp1[child] > firstMax) {
secondMax = firstMax;
firstMax = dp1[child];
}
else if (dp1[child] > secondMax) {
secondMax = dp1[child];
}
}
// Calculate dp1[node] and dp2[node]
dp1[node] = 1 + (firstMax == -1 ? 0 : firstMax);
dp2[node] = (firstMax == -1 ? 0 : firstMax) +
(secondMax == -1 ? 0 : secondMax);
// Update diameter
diameter = max(diameter, dp2[node]);
return dp1[node];}
// Function to calculate the diameter of the tree int diameter(Node* root) { if (!root) return 0;
unordered_map<Node*, int> dp1, dp2;
int diameter = 0;
// Call DFS to calculate diameter
dfs(root, dp1, dp2, diameter);
return diameter;}
int main() {
// Hard coded tree
// 1
// / \
// 2 3
// / \
// 4 5
Node* root = new Node(1);
Node* node2 = new Node(2);
Node* node3 = new Node(3);
Node* node4 = new Node(4);
Node* node5 = new Node(5);
root->children.push_back(node2);
root->children.push_back(node3);
node2->children.push_back(node4);
node2->children.push_back(node5);
cout << diameter(root) << endl;
return 0;}
Java
// Java program to find the diameter // of n-ary tree using dp. import java.util.*;
class Node { int data; List children;
Node(int x) {
data = x;
children = new ArrayList<>();
}}
class GfG {
// Function to perform DFS and calculate dp1 and dp2
static int dfs(Node node, Map<Node, Integer> dp1,
Map<Node, Integer> dp2, int[] diameter) {
if (node == null) return 0;
int firstMax = -1, secondMax = -1;
// Traverse children of the current node
for (Node child : node.children) {
// Recursive call to process child
dfs(child, dp1, dp2, diameter);
// Find first and second maximum dp1 values
if (dp1.getOrDefault(child, 0) > firstMax) {
secondMax = firstMax;
firstMax = dp1.getOrDefault(child, 0);
}
else if (dp1.getOrDefault(child, 0) > secondMax) {
secondMax = dp1.getOrDefault(child, 0);
}
}
// Calculate dp1[node] and dp2[node]
dp1.put(node, 1 + (firstMax == -1 ? 0 : firstMax));
dp2.put(node, (firstMax == -1 ? 0 : firstMax) +
(secondMax == -1 ? 0 : secondMax));
// Update diameter
diameter[0] = Math.max(diameter[0],
dp2.getOrDefault(node, 0));
return dp1.get(node);
}
// Function to calculate the diameter of the tree
static int diameter(Node root) {
if (root == null) return 0;
Map<Node, Integer> dp1 = new HashMap<>();
Map<Node, Integer> dp2 = new HashMap<>();
int[] diameter = new int[1];
// Call DFS to calculate diameter
dfs(root, dp1, dp2, diameter);
return diameter[0];
}
public static void main(String[] args) {
// Hard coded tree
// 1
// / \
// 2 3
// / \
// 4 5
Node root = new Node(1);
Node node2 = new Node(2);
Node node3 = new Node(3);
Node node4 = new Node(4);
Node node5 = new Node(5);
root.children.add(node2);
root.children.add(node3);
node2.children.add(node4);
node2.children.add(node5);
System.out.println(diameter(root));
}}
Python
Python program to find the diameter
of n-ary tree using dp.
class Node: def init(self, x): self.data = x self.children = []
Function to perform DFS and calculate dp1 and dp2
def dfs(node, dp1, dp2, diameter): if not node: return 0
firstMax = -1
secondMax = -1
# Traverse children of the current node
for child in node.children:
# Recursive call to process child
dfs(child, dp1, dp2, diameter)
# Find first and second maximum dp1 values
if dp1.get(child, 0) > firstMax:
secondMax = firstMax
firstMax = dp1.get(child, 0)
elif dp1.get(child, 0) > secondMax:
secondMax = dp1.get(child, 0)
# Calculate dp1[node] and dp2[node]
dp1[node] = 1 + (0 if firstMax == -1 else firstMax)
dp2[node] = (0 if firstMax == -1 else firstMax) + \
(0 if secondMax == -1 else secondMax)
# Update diameter
diameter[0] = max(diameter[0], dp2[node])
return dp1[node]Function to calculate the diameter of the tree
def diameter(root): if not root: return 0
dp1 = {}
dp2 = {}
diameter = [0]
# Call DFS to calculate diameter
dfs(root, dp1, dp2, diameter)
return diameter[0]if name == "main":
# Hard coded tree
# 1
# / \
# 2 3
# / \
# 4 5
root = Node(1)
node2 = Node(2)
node3 = Node(3)
node4 = Node(4)
node5 = Node(5)
root.children.extend([node2, node3])
node2.children.extend([node4, node5])
print(diameter(root))C#
// C# program to find the diameter // of n-ary tree using dp. using System; using System.Collections.Generic;
class Node { public int data; public List children;
public Node(int x) {
data = x;
children = new List<Node>();
}}
class GfG {
// Function to perform DFS and calculate dp1 and dp2
static int dfs(Node node, Dictionary<Node, int> dp1,
Dictionary<Node, int> dp2, ref int diameter) {
if (node == null) return 0;
int firstMax = -1, secondMax = -1;
// Traverse children of the current node
foreach (Node child in node.children) {
// Recursive call to process child
dfs(child, dp1, dp2, ref diameter);
// Find first and second maximum dp1 values
if (dp1.ContainsKey(child) && dp1[child] > firstMax) {
secondMax = firstMax;
firstMax = dp1[child];
} else if (dp1.ContainsKey(child) && dp1[child] > secondMax) {
secondMax = dp1[child];
}
}
// Calculate dp1[node] and dp2[node]
dp1[node] = 1 + (firstMax == -1 ? 0 : firstMax);
dp2[node] = (firstMax == -1 ? 0 : firstMax) +
(secondMax == -1 ? 0 : secondMax);
// Update diameter
diameter = Math.Max(diameter, dp2[node]);
return dp1[node];
}
// Function to calculate the diameter of the tree
static int diameter(Node root) {
if (root == null) return 0;
Dictionary<Node, int> dp1 = new Dictionary<Node, int>();
Dictionary<Node, int> dp2 = new Dictionary<Node, int>();
int diameter = 0;
// Call DFS to calculate diameter
dfs(root, dp1, dp2, ref diameter);
return diameter;
}
static void Main(string[] args) {
// Hard coded tree
// 1
// / \
// 2 3
// / \
// 4 5
Node root = new Node(1);
Node node2 = new Node(2);
Node node3 = new Node(3);
Node node4 = new Node(4);
Node node5 = new Node(5);
root.children.Add(node2);
root.children.Add(node3);
node2.children.Add(node4);
node2.children.Add(node5);
Console.WriteLine(diameter(root));
}}
JavaScript
// JavaScript program to find the diameter // of n-ary tree using dp.
class Node { constructor(x) { this.data = x; this.children = []; } }
// Function to perform DFS and calculate dp1 and dp2 function dfs(node, dp1, dp2, diameter) { if (!node) return 0;
let firstMax = -1, secondMax = -1;
// Traverse children of the current node
for (const child of node.children) {
// Recursive call to process child
dfs(child, dp1, dp2, diameter);
// Find first and second maximum dp1 values
if ((dp1[child] || 0) > firstMax) {
secondMax = firstMax;
firstMax = dp1[child] || 0;
} else if ((dp1[child] || 0) > secondMax) {
secondMax = dp1[child] || 0;
}
}
// Calculate dp1[node] and dp2[node]
dp1[node] = 1 + (firstMax === -1 ? 0 : firstMax);
dp2[node] = (firstMax === -1 ? 0 : firstMax) +
(secondMax === -1 ? 0 : secondMax);
// Update diameter
diameter[0] = Math.max(diameter[0], dp2[node]);
return dp1[node];}
// Function to calculate the diameter of the tree function diameter(root) { if (!root) return 0;
let dp1 = {};
let dp2 = {};
let diameter = [0];
// Call DFS to calculate diameter
dfs(root, dp1, dp2, diameter);
return diameter[0];}
// Hard coded tree
// 1
// /
// 2 3
// /
// 4 5
const root = new Node(1);
const node2 = new Node(2);
const node3 = new Node(3);
const node4 = new Node(4);
const node5 = new Node(5);
root.children.push(node2, node3);
node2.children.push(node4, node5);
console.log(diameter(root));
`
**Time Complexity: O(n), where n is the total number of nodes in the tree.
**Auxiliary Space: O(n)
**Related articles: