1522. Diameter of N-Ary Tree


Problem Description

The problem involves finding the diameter of an N-ary tree. The diameter is defined as the length (number of edges) of the longest path between any two nodes within the tree. A key point to note is that this path is not required to go through the tree's root. Understanding how to handle N-ary trees, which are trees where a node can have any number of children, is crucial in solving this problem. Level order traversal serialization (where each group of children is separated by the null value) is used for input representation.

Intuition

When finding the diameter of an N-ary tree, a depth-first search (DFS) approach can be applied effectively. The diameter of the tree could potentially be the distance between two nodes that are below the root of the tree, which means it might not include the root itself. Therefore, for each node, we need to consider the two longest paths that extend from this node down its children, because these paths could potentially contribute to the maximum diameter if they are connected through the current node.

To implement this, we use a recursive DFS function that computes and returns the height of the current subtree while simultaneously calculating the potential diameter that passes through the current node (the sum of the two longest child paths connected at the node). We track the overall longest path seen so far with a nonlocal variable 'ans', which gets updated when a larger diameter is found.

The key steps in our approach are:

  1. Traverse each node with DFS, starting from the root.
  2. For each node visited, calculate the lengths of the longest (max) and second-longest (smax) paths among its child nodes.
  3. Update the overall potential diameter 'ans' if the sum of max and smax is greater than the current 'ans'.
  4. At each node, return the height of this subtree to its parent, which is 1 plus the longest path length among the child nodes.

By doing this for each node, once the traversal is complete, 'ans' will hold the length of the longest path, which is the diameter of the N-ary tree.

Learn more about Tree and Depth-First Search patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

Given a sorted array of integers and an integer called target, find the element that equals to the target and return its index. Select the correct code that fills the ___ in the given code snippet.

1def binary_search(arr, target):
2    left, right = 0, len(arr) - 1
3    while left ___ right:
4        mid = (left + right) // 2
5        if arr[mid] == target:
6            return mid
7        if arr[mid] < target:
8            ___ = mid + 1
9        else:
10            ___ = mid - 1
11    return -1
12
1public static int binarySearch(int[] arr, int target) {
2    int left = 0;
3    int right = arr.length - 1;
4
5    while (left ___ right) {
6        int mid = left + (right - left) / 2;
7        if (arr[mid] == target) return mid;
8        if (arr[mid] < target) {
9            ___ = mid + 1;
10        } else {
11            ___ = mid - 1;
12        }
13    }
14    return -1;
15}
16
1function binarySearch(arr, target) {
2    let left = 0;
3    let right = arr.length - 1;
4
5    while (left ___ right) {
6        let mid = left + Math.trunc((right - left) / 2);
7        if (arr[mid] == target) return mid;
8        if (arr[mid] < target) {
9            ___ = mid + 1;
10        } else {
11            ___ = mid - 1;
12        }
13    }
14    return -1;
15}
16

Solution Approach

The solution approach for calculating the diameter of an N-ary tree involves a depth-first search (DFS) algorithm. The DFS is customized to perform two operations at once – calculate the height of the current subtree and use this information to evaluate the longest path that passes through each node (potential diameters).

Here's a step-by-step breakdown of the implementation using the provided solution code:

  1. We define a DFS helper function dfs that takes a node of the tree as an argument. This function returns the height of the tree rooted at the given node.

  2. Inside the dfs function:

    • We start by checking if the current root node is None (base case). If it is, we return 0 since a non-existent node does not contribute to height or diameter.

    • We then introduce two variables m1 and m2 initialized to 0, which will store the lengths of the longest (m1) and second-longest (m2) paths found within the children of the current node.

    • We loop through each child of the current node and recursively call dfs(child). The returned value t is the height of the subtree rooted at child. We use this value to potentially update m1 and m2.

      • If t is greater than m1, we set m2 to m1 and then m1 to t, thus keeping m1 as the max height and m2 as the second max height among the children.

      • If t is not greater than m1 but is greater than m2, we update m2 with t.

    • After considering all children, we update the ans variable (which is kept in the outer scope of dfs and declared as nonlocal) with the maximum of its current value and the sum of m1 + m2, representing the potential diameter at this node (the longest path through the node).

  3. The dfs function concludes by returning 1 + m1, which represents the height of the subtree rooted at the current node (1 for the edge connecting to its parent and m1 as the height of its longest subtree).

  4. In the diameter method of the Solution class, we initialize the ans variable to 0 (to globally keep track of the longest path seen) and call dfs(root) to kick off the recursive DFS from the root of the tree.

  5. Finally, the diameter method returns the ans variable, which now contains the diameter of the tree.

The use of recursion to traverse the tree, combined with updating the two longest paths at each node, harnesses the depth-first search pattern efficiently to solve the problem of finding the tree's diameter.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

What's the output of running the following function using input 56?

1KEYBOARD = {
2    '2': 'abc',
3    '3': 'def',
4    '4': 'ghi',
5    '5': 'jkl',
6    '6': 'mno',
7    '7': 'pqrs',
8    '8': 'tuv',
9    '9': 'wxyz',
10}
11
12def letter_combinations_of_phone_number(digits):
13    def dfs(path, res):
14        if len(path) == len(digits):
15            res.append(''.join(path))
16            return
17
18        next_number = digits[len(path)]
19        for letter in KEYBOARD[next_number]:
20            path.append(letter)
21            dfs(path, res)
22            path.pop()
23
24    res = []
25    dfs([], res)
26    return res
27
1private static final Map<Character, char[]> KEYBOARD = Map.of(
2    '2', "abc".toCharArray(),
3    '3', "def".toCharArray(),
4    '4', "ghi".toCharArray(),
5    '5', "jkl".toCharArray(),
6    '6', "mno".toCharArray(),
7    '7', "pqrs".toCharArray(),
8    '8', "tuv".toCharArray(),
9    '9', "wxyz".toCharArray()
10);
11
12public static List<String> letterCombinationsOfPhoneNumber(String digits) {
13    List<String> res = new ArrayList<>();
14    dfs(new StringBuilder(), res, digits.toCharArray());
15    return res;
16}
17
18private static void dfs(StringBuilder path, List<String> res, char[] digits) {
19    if (path.length() == digits.length) {
20        res.add(path.toString());
21        return;
22    }
23    char next_digit = digits[path.length()];
24    for (char letter : KEYBOARD.get(next_digit)) {
25        path.append(letter);
26        dfs(path, res, digits);
27        path.deleteCharAt(path.length() - 1);
28    }
29}
30
1const KEYBOARD = {
2    '2': 'abc',
3    '3': 'def',
4    '4': 'ghi',
5    '5': 'jkl',
6    '6': 'mno',
7    '7': 'pqrs',
8    '8': 'tuv',
9    '9': 'wxyz',
10}
11
12function letter_combinations_of_phone_number(digits) {
13    let res = [];
14    dfs(digits, [], res);
15    return res;
16}
17
18function dfs(digits, path, res) {
19    if (path.length === digits.length) {
20        res.push(path.join(''));
21        return;
22    }
23    let next_number = digits.charAt(path.length);
24    for (let letter of KEYBOARD[next_number]) {
25        path.push(letter);
26        dfs(digits, path, res);
27        path.pop();
28    }
29}
30

Example Walkthrough

Let's illustrate the solution approach with a small example. Suppose we have an N-ary tree represented using level order serialization as follows:

1     1
2   / | \
3  2  3  4
4 /      | \
55       6  7

We want to find the diameter of this tree, which is the longest path between any two nodes. We will walk through the DFS approach using the provided solution template:

  1. We start the DFS with the root, which is node 1. We initialize ans to 0. The dfs function will traverse the tree and track the longest path at each node.

  2. When we process the root (node 1), we initialize m1 and m2 both to 0. These variables will track the longest and second-longest heights among the children of each node.

  3. We recursively call dfs for each child of the root.

    • For child node 2, the recursive call to dfs returns 1 (since node 5 is its only child and adds one edge to the height).
    • For child node 3, the recursive call to dfs returns 0 (as it has no child).
    • For child node 4, the DFS must consider both children. Nodes 6 and 7 contribute a height of 1 each, and since these are the longest and second-longest heights for this subtree, after considering both, node 4 returns 2 as the height (1 for the longest child plus 1 for the edge to node 4).
  4. Now, with the heights from the children of the root, we update m1 and m2 for the root. m1 becomes 2 (from node 4), and m2 becomes 1 (from node 2). Therefore, the potential diameter passing through root at this stage is m1 + m2 = 2 + 1 = 3. We update ans to 3.

  5. The DFS continues recursively for all nodes but finds no longer path than the current ans.

  6. After the DFS is complete and all nodes have been visited, ans holds the diameter of the tree. In this case, it remains 3, representing the length of the longest path, which happens to be from node 5 to node 7 via the root (5 -> 2 -> 1 -> 4 -> 6 or 7).

  7. The diameter method then returns ans, which is 3, as the diameter of the tree.

This walkthrough illustrates the algorithm's efficiency in computing the diameter by determining the longest path through each node without necessarily passing through the root.

Solution Implementation

1# Definition for a Node.
2class Node:
3    def __init__(self, val=None, children=None):
4        self.val = val
5        self.children = children if children is not None else []
6
7class Solution:
8    def diameter(self, root: 'Node') -> int:
9        # Helper function to perform depth-first search
10        def dfs(node: 'Node') -> int:
11            # base case: if the current node is None, return 0
12            if node is None:
13                return 0
14            # accessing the non-local variable 'max_diameter' to update its value within this helper function
15            nonlocal max_diameter
16            # initialize the two longest paths from the current node to 0
17            longest_path = second_longest_path = 0
18            # iterate over all the children of the current node
19            for child in node.children:
20                # recursively find the longest path for each child
21                path_length = dfs(child)
22                # check if the current path is longer than the longest recorded path
23                if path_length > longest_path:
24                    # update the second longest and longest paths accordingly
25                    second_longest_path, longest_path = longest_path, path_length
26                # else if it's only longer than the second longest, update the second longest
27                elif path_length > second_longest_path:
28                    second_longest_path = path_length
29            # update the maximum diameter encountered so far based on the two longest paths from this node
30            max_diameter = max(max_diameter, longest_path + second_longest_path)
31            # return the longer path increased by 1 for the edge between this node and its parent
32            return 1 + longest_path
33
34        # Initialize max_diameter to 0 before starting DFS
35        max_diameter = 0
36        # Call the dfs function starting from the root node
37        dfs(root)
38        # Once DFS is complete, return the max_diameter found
39        return max_diameter
40
1/**
2 * Definition for a N-ary tree node.
3 */
4class Node {
5    public int val;
6    public List<Node> children;
7
8    // Constructor initializes value and empty children list
9    public Node() {
10        children = new ArrayList<Node>();
11    }
12
13    // Constructor initializes node with a value and empty children list
14    public Node(int val) {
15        this.val = val;
16        children = new ArrayList<Node>();
17    }
18
19    // Constructor initializes node with a value and given children list
20    public Node(int val, ArrayList<Node> children) {
21        this.val = val;
22        this.children = children;
23    }
24}
25
26public class Solution {
27    // This class variable keeps track of the diameter of the tree
28    private int diameter;
29
30    /**
31     * Computes the diameter of an N-ary tree.
32     * The diameter of an N-ary tree is the length of the longest path between any two nodes in a tree.
33     *
34     * @param root the root node of the tree
35     * @return the diameter of the tree
36     */
37    public int diameter(Node root) {
38        diameter = 0;
39        depthFirstSearch(root);
40        return diameter;
41    }
42
43    /**
44     * Uses Depth First Search (DFS) to find the length of the longest path through the N-ary tree from the current node.
45     *
46     * @param root the current node being traversed
47     * @return the maximum depth emanating from the current node
48     */
49    private int depthFirstSearch(Node root) {
50        // Leaf nodes have a depth of 0
51        if (root == null) {
52            return 0;
53        }
54
55        int maxDepth1 = 0; // Tracks the longest path
56        int maxDepth2 = 0; // Tracks the second longest path
57
58        // Recursively obtain the depth for children nodes to find the longest paths
59        for (Node child : root.children) {
60            int depth = depthFirstSearch(child);
61          
62            // Check and set the longest and second-longest distances found
63            if (depth > maxDepth1) {
64                maxDepth2 = maxDepth1;
65                maxDepth1 = depth;
66            } else if (depth > maxDepth2) {
67                maxDepth2 = depth;
68            }
69        }
70
71        // Update the diameter if the sum of two longest paths through the current node is greater than the current diameter
72        diameter = Math.max(diameter, maxDepth1 + maxDepth2);
73
74        // Return the maximum depth plus one for the current path
75        return 1 + maxDepth1;
76    }
77}
78
1// Definition for a Node is provided as per the question context
2class Node {
3public:
4    int val;                      // The value contained within the node.
5    vector<Node*> children;       // A vector containing pointers to its children.
6
7    Node() {}
8
9    Node(int _val) {
10        val = _val;
11    }
12
13    Node(int _val, vector<Node*> _children) {
14        val = _val;
15        children = _children;
16    }
17};
18
19class Solution {
20public:
21    int maxDiameter;  // Class variable to store the maximum diameter found.
22
23    // Public method which is the starting point for finding the diameter of the tree.
24    int diameter(Node* root) {
25        maxDiameter = 0; // Initialize the max diameter to 0.
26        dfs(root);       // Call the depth-first search helper method.
27        return maxDiameter; // Return the maximum diameter calculated.
28    }
29
30    // Private helper method for DFS traversal which calculates the depths and updates the diameter.
31    // It returns the maximum depth from the current node to its furthest leaf node.
32    int dfs(Node* root) {
33        if (!root) return 0; // Base case: If the node is null, return 0 (no depth).
34      
35        int maxDepth1 = 0; // To store the maximum length of the paths in the children.
36        int maxDepth2 = 0; // To store the second maximum length of the paths in the children.
37
38        // Iterate through each child node of the current root.
39        for (Node* child : root->children) {
40            int depth = dfs(child); // Recursive call to get the depth for each child.
41            // Update the two maximum depths found among children nodes
42            if (depth > maxDepth1) {
43                maxDepth2 = maxDepth1; // Update the second max if the new max is found
44                maxDepth1 = depth;    // Update the new max depth
45            } else if (depth > maxDepth2) {
46                maxDepth2 = depth; // Update the second max if greater than it but less than the max depth
47            }
48        }
49
50        // Update the maximum diameter if the sum of the two largest depths is greater than the current diameter.
51        maxDiameter = max(maxDiameter, maxDepth1 + maxDepth2);
52
53        // Return the maximum depth of this subtree to its parent caller, which is 1 plus the max depth of its children.
54        return 1 + maxDepth1;
55    }
56};
57
1// Define the type for a Node that includes a value and an array of Node references as children
2type Node = {
3  val: number;
4  children: Node[];
5};
6
7let maxDiameter: number; // Global variable to store the maximum diameter found.
8
9// This function is the entry point for finding the diameter of the n-ary tree.
10function diameter(root: Node | null): number {
11  maxDiameter = 0; // Initialize the max diameter to 0.
12  dfs(root); // Call the depth-first search helper function.
13  return maxDiameter; // Return the maximum diameter calculated.
14}
15
16// Helper function for DFS traversal that calculates depths and updates the diameter.
17// It returns the maximum distance from the current node to its furthest leaf node.
18function dfs(root: Node | null): number {
19  if (root === null) return 0; // Base case: If the node is null, return 0 (no depth).
20
21  let maxDepthOne = 0; // To store the maximum depth among the paths in the children.
22  let maxDepthTwo = 0; // To store the second maximum depth among the paths in the children.
23
24  // Iterate through each child node of the current root.
25  for (let child of root.children) {
26    let depth = dfs(child); // Recursive call to get the depth for each child.
27    // Update the two maximum depths found among children nodes
28    if (depth > maxDepthOne) {
29      maxDepthTwo = maxDepthOne; // Second max updated if a new max is found
30      maxDepthOne = depth; // Set as the new max depth
31    } else if (depth > maxDepthTwo) {
32      maxDepthTwo = depth; // Update the second max if it's greater than the current second max
33    }
34  }
35
36  // Update the maximum diameter if the sum of the two largest depths is greater than the current maximum diameter.
37  maxDiameter = Math.max(maxDiameter, maxDepthOne + maxDepthTwo);
38
39  // Return the maximum depth of this subtree to its parent, which is 1 plus the maximum depth of its children.
40  return 1 + maxDepthOne;
41}
42
Not Sure What to Study? Take the 2-min Quiz:

Which of the following is the prefix sum of array [1, 2, 3, 4, 5]?

Time and Space Complexity

The code provided defines a function dfs(root) that is used to compute the diameter of an N-ary tree. The diameter of an N-ary tree is defined as the length of the longest path between any two nodes in the tree.

Time Complexity:

The time complexity of the provided code is O(N), where N is the total number of nodes in the tree. This is because the dfs() function is called recursively for each node exactly once. Within each call to dfs(), it performs a constant amount of work for each child. Since the sum of the sizes of all children's lists across all nodes is N - 1 (there are N - 1 edges in a tree with N nodes), the overall time complexity is linear relative to the number of nodes in the tree.

Space Complexity:

The space complexity of the code can be analyzed in two parts: the space required for the recursion call stack and the space required for the dfs() function execution.

  • The worst-case space complexity for the recursion call stack is O(H), where H is the height of the tree. In the worst case, the tree can be skewed, resembling a linked list, thus the height can become N, leading to a worst-case space complexity of O(N).

  • The additional space used by the dfs() function is constant, as it only uses a few integer variables (m1, m2, and t).

Considering both parts, the total space complexity is O(H), which is O(N) in the worst case (when the tree is a linear chain), but can be better (such as O(log N)) if the tree is balanced.

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

Which of the following is equvalent to O(3*2^n + n^3 + n!+ log n)?


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫