Problem Description

You are given the root of an N-ary tree. An N-ary tree is a tree in which each node can have zero or more children (not limited to two as in binary trees). Your task is to compute the diameter of the tree.

The diameter of an N-ary tree is defined as the length of the longest path between any two nodes in the tree. The path may or may not pass through the root. The length of a path is measured as the number of edges between the two nodes.

• Each node has a value and a list of child nodes.
• The tree is rooted at a given node root.
• You must return an integer representing the diameter (the number of edges in the longest path between any two nodes).

Constraints:

• The number of nodes in the tree is in the range [1, 10^4].
• The tree is well-formed (no cycles).
• Each node's children list is not null (may be empty).

Thought Process

To solve this problem, let's first understand what the diameter means in the context of an N-ary tree. Unlike binary trees, each node can have many children, so the branching factor is not fixed. The diameter is the longest path (in terms of edges) between any two nodes, which may or may not pass through the root.

A brute-force approach would be to compute the longest path between every pair of nodes, but this would be extremely inefficient for large trees. Instead, we can use a recursive approach similar to how we find the diameter in binary trees, but generalized for N-ary trees.

The key insight is that the diameter at a node can be the sum of the two largest heights among its children, because the longest path passing through that node would go down to the deepest two children (possibly different subtrees). So, for each node, we want to track the two largest subtree heights among its children.

Solution Approach

We use a depth-first search (DFS) traversal to process the tree from the bottom up. For each node, we:

1. Recursively compute the height of each child subtree.
2. Track the two largest heights among all children. Let's call them max1 and max2.
3. Update the global diameter as max1 + max2 if it's greater than the current best.
4. Return max1 + 1 as the height of the current node (since we add one edge to reach the child).

This way, we ensure that at every node, we consider the possibility that the longest path may pass through that node, connecting its two deepest subtrees. We use a variable (often nonlocal or class-level) to keep track of the maximum diameter found so far.

• Why DFS? Because we need to know the heights of all child subtrees before we can compute the diameter at a node.
• Why two largest heights? Because the longest path passing through a node will go down the two deepest subtrees.
• Why not more than two? Any path in a tree cannot fork; it must be a single path, so at most two branches can contribute, one going down each direction.

Example Walkthrough

Consider the following N-ary tree:

        1
      / | \
     2  3  4
        |   \
        5    6
  

• Node 1 has children 2, 3, 4.
• Node 3 has child 5.
• Node 4 has child 6.

Let's compute the heights:

• Nodes 2, 5, 6 are leaves: height = 0.
• Node 3: only child is 5 (height 0) ⇒ height = 1.
• Node 4: only child is 6 (height 0) ⇒ height = 1.
• Node 1: children heights = [0 (2), 1 (3), 1 (4)]. The two largest are 1 and 1.

• From 5 to 1: 5–3–1 (2 edges)
• From 6 to 1: 6–4–1 (2 edges)
• From 5 to 6: 5–3–1–4–6 (4 edges) ← This is the diameter

Time and Space Complexity

Brute-force approach:

• For each pair of nodes, compute the distance between them (using BFS/DFS). There are O(N^2) pairs, and each distance can take O(N) time, so total time is O(N^3), which is infeasible for large N.

• Each node is visited once, and for each node, we process all its children. Total time is O(N), where N is the number of nodes.
• Space complexity is O(H), where H is the height of the tree, due to recursion stack. In the worst case (skewed tree), H = N, so space is O(N).

Summary

The diameter of an N-ary tree can be efficiently computed using a single DFS traversal. By keeping track of the two largest subtree heights at each node, we can update the maximum possible diameter as we traverse the tree. This approach is elegant because it generalizes the binary tree diameter method, operates in linear time, and uses only a small amount of extra space. Understanding how the diameter relates to the sum of the two largest child heights at each node is the key insight.

Code Implementation

# Definition for a Node.
class Node:
    def __init__(self, val=None, children=None):
        self.val = val
        self.children = children if children is not None else []

class Solution:
    def diameter(self, root: 'Node') -> int:
        self.max_diameter = 0

        def dfs(node):
            if not node:
                return 0
            max1, max2 = 0, 0  # The two largest heights among children
            for child in node.children:
                h = dfs(child)
                if h > max1:
                    max2 = max1
                    max1 = h
                elif h > max2:
                    max2 = h
            # Update the diameter at this node
            self.max_diameter = max(self.max_diameter, max1 + max2)
            return max1 + 1  # Height of this node

        dfs(root)
        return self.max_diameter
      

/*
// Definition for a Node.
class Node {
public:
    int val;
    vector<Node*> children;

    Node() {}

    Node(int _val) {
        val = _val;
    }

    Node(int _val, vector<Node*> _children) {
        val = _val;
        children = _children;
    }
};
*/

class Solution {
public:
    int maxDiameter = 0;

    int dfs(Node* node) {
        if (!node) return 0;
        int max1 = 0, max2 = 0;
        for (Node* child : node->children) {
            int h = dfs(child);
            if (h > max1) {
                max2 = max1;
                max1 = h;
            } else if (h > max2) {
                max2 = h;
            }
        }
        maxDiameter = max(maxDiameter, max1 + max2);
        return max1 + 1;
    }

    int diameter(Node* root) {
        dfs(root);
        return maxDiameter;
    }
};
      

/*
// Definition for a Node.
class Node {
    public int val;
    public List<Node> children;

    public Node() {}

    public Node(int _val) {
        val = _val;
    }

    public Node(int _val, List<Node> _children) {
        val = _val;
        children = _children;
    }
};
*/

class Solution {
    private int maxDiameter = 0;

    private int dfs(Node node) {
        if (node == null) return 0;
        int max1 = 0, max2 = 0;
        for (Node child : node.children) {
            int h = dfs(child);
            if (h > max1) {
                max2 = max1;
                max1 = h;
            } else if (h > max2) {
                max2 = h;
            }
        }
        maxDiameter = Math.max(maxDiameter, max1 + max2);
        return max1 + 1;
    }

    public int diameter(Node root) {
        dfs(root);
        return maxDiameter;
    }
}
      

/**
 * // Definition for a Node.
 * function Node(val, children) {
 *    this.val = val;
 *    this.children = children || [];
 * };
 */

/**
 * @param {Node} root
 * @return {number}
 */
var diameter = function(root) {
    let maxDiameter = 0;

    function dfs(node) {
        if (!node) return 0;
        let max1 = 0, max2 = 0;
        for (let child of node.children) {
            let h = dfs(child);
            if (h > max1) {
                max2 = max1;
                max1 = h;
            } else if (h > max2) {
                max2 = h;
            }
        }
        maxDiameter = Math.max(maxDiameter, max1 + max2);
        return max1 + 1;
    }

    dfs(root);
    return maxDiameter;
};