DP on Trees | Set-3 ( Diameter of N-ary Tree )

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<Node> 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# program to find the diameter
// of n-ary tree using dp.
using System;
using System.Collections.Generic;

class Node {
    public int data;
    public List<Node> 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));

Output

Time Complexity: O(n), where n is the total number of nodes in the tree.
Auxiliary Space: O(n)

Related articles:

DP on Trees | Set-3 ( Diameter of N-ary Tree )

Explore