LeetCode 1522 - Diameter of N-Ary Tree

This problem asks us to compute the diameter of an N-ary tree. An N-ary tree is a tree where each node can have any numb

LeetCode Problem 1522

Difficulty: 🟡 Medium
Topics: Tree, Depth-First Search

Solution

LeetCode 1522 - Diameter of N-Ary Tree

Problem Understanding

This problem asks us to compute the diameter of an N-ary tree. An N-ary tree is a tree where each node can have any number of children, unlike a binary tree where each node has at most two children.

The diameter of a tree is defined as the length of the longest path between any two nodes in the tree. The path does not need to pass through the root. The length of a path is measured by the number of edges between the two nodes.

For example, if the longest path connects nodes through 3 edges, then the diameter is 3.

The input is provided as the root node of the N-ary tree. Each node contains:

  • A value
  • A list of child nodes

The serialization format shown in the examples uses level-order traversal with null separators between groups of children, but in the actual solution we only interact with the tree structure itself.

The constraints are important:

  • The tree depth is at most 1000
  • The number of nodes is at most 10^4

These limits tell us that an O(n^2) solution could become too slow for large trees. We should aim for a linear time solution, ideally O(n).

Several edge cases are important:

  • A tree with only one node has diameter 0, because there are no edges.
  • A highly skewed tree behaves like a linked list.
  • The longest path may not pass through the root.
  • Nodes may have many children, so we cannot rely on binary-tree-specific logic.

The key challenge is determining the longest path efficiently while visiting each node only once.

Approaches

Brute Force Approach

A brute force solution would compute the longest path between every pair of nodes.

One way to do this is:

  1. Treat the tree as an undirected graph.
  2. For every node:
  • Run DFS or BFS to compute distances to all other nodes.
  • Track the maximum distance found.

This works because the diameter is simply the largest shortest-path distance between any two nodes.

However, this approach is very inefficient. If there are n nodes, we perform a traversal from every node, resulting in O(n^2) time complexity.

With up to 10^4 nodes, this becomes too expensive.

Optimal Approach

The key insight is that the diameter can be computed during a single depth-first traversal.

For every node, we want to know:

  • The longest downward path starting from that node.
  • Whether combining two child paths through the current node creates a larger diameter.

This is similar to the diameter problem in binary trees, except here each node may have many children.

At each node:

  • Compute the height of every child subtree.
  • Keep track of the two largest child heights.
  • The longest path passing through the current node is:
largest_height + second_largest_height

This works because the longest path through a node must connect two deepest child branches.

While recursively computing heights, we continuously update a global diameter value.

Since each node is processed exactly once, the solution runs in linear time.

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(n) Run traversal from every node
Optimal O(n) O(h) Single DFS traversal using subtree heights

Here, h is the height of the tree.

Algorithm Walkthrough

  1. Initialize a global variable diameter to 0.

This variable stores the largest diameter found so far while traversing the tree. 2. Perform a depth-first search starting from the root.

The DFS function will return the height of the subtree rooted at the current node. 3. For each node, examine all of its children.

For every child:

  • Recursively compute the child's subtree height.
  • Add 1 to represent the edge from the current node to the child.
  1. Track the two largest child heights.

We only need the top two heights because the longest path through the current node can only use two branches. 5. Compute the candidate diameter through the current node.

The path length is:

longest_height + second_longest_height

Update the global diameter if this value is larger. 6. Return the largest downward height to the parent.

The parent only needs the single longest downward branch. 7. Continue recursively until all nodes are processed. 8. Return the final diameter.

Why it works

For every node, the longest path that passes through that node must connect the two deepest child subtrees. By computing subtree heights bottom-up, we guarantee that when processing a node, all child heights are already known.

Since every possible diameter either:

  • Lies entirely inside a subtree, or
  • Passes through some node connecting two child branches,

checking every node guarantees that the true diameter is found.

Python Solution

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

from typing import Optional, List

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

        def dfs(node: 'Node') -> int:
            if not node:
                return 0

            longest = 0
            second_longest = 0

            for child in node.children:
                child_height = dfs(child) + 1

                if child_height > longest:
                    second_longest = longest
                    longest = child_height
                elif child_height > second_longest:
                    second_longest = child_height

            self.max_diameter = max(
                self.max_diameter,
                longest + second_longest
            )

            return longest

        dfs(root)

        return self.max_diameter

The solution uses recursive DFS to compute subtree heights.

The dfs() function returns the height of the subtree rooted at the current node. While processing children, the code tracks the two largest child heights using longest and second_longest.

Whenever a larger child height is found:

  • The previous largest becomes the second largest.
  • The new height becomes the largest.

At every node, the potential diameter through that node is:

longest + second_longest

The global maximum diameter is updated accordingly.

Finally, the function returns the longest downward path so the parent node can continue building larger paths.

Go Solution

/**
 * Definition for a Node.
 * type Node struct {
 *     Val int
 *     Children []*Node
 * }
 */

func diameter(root *Node) int {
    maxDiameter := 0

    var dfs func(node *Node) int

    dfs = func(node *Node) int {
        if node == nil {
            return 0
        }

        longest := 0
        secondLongest := 0

        for _, child := range node.Children {
            childHeight := dfs(child) + 1

            if childHeight > longest {
                secondLongest = longest
                longest = childHeight
            } else if childHeight > secondLongest {
                secondLongest = childHeight
            }
        }

        if longest+secondLongest > maxDiameter {
            maxDiameter = longest + secondLongest
        }

        return longest
    }

    dfs(root)

    return maxDiameter
}

The Go implementation follows the same logic as the Python solution.

A closure-based DFS function is used so it can access the shared maxDiameter variable.

Go uses nil checks instead of Python's truthy checks. Child nodes are stored as slices, and iteration is performed using a range loop.

Since the constraints are small enough for standard integer operations, integer overflow is not a concern.

Worked Examples

Example 1

Input:

[1,null,3,2,4,null,5,6]

Tree structure:

        1
      / | \
     3  2  4
    / \
   5   6

The longest path is:

5 -> 3 -> 1 -> 2

This path contains 3 edges.

DFS Processing Table

Node Child Heights Longest Second Longest Diameter Update
5 none 0 0 0
6 none 0 0 0
3 1, 1 1 1 2
2 none 0 0 2
4 none 0 0 2
1 2, 1, 1 2 1 3

Final answer:

3

Example 2

Input:

[1,null,2,null,3,4,null,5,null,6]

Tree structure:

1
|
2
|
3
/ \
5   4
|
6

Longest path:

6 -> 5 -> 3 -> 2 -> 1

This path contains 4 edges.

DFS Processing Table

Node Child Heights Longest Second Longest Diameter Update
6 none 0 0 0
5 1 1 0 1
4 none 0 0 1
3 2, 1 2 1 3
2 3 3 0 3
1 4 4 0 4

Final answer:

4

Example 3

The deepest branches are connected through internal nodes rather than the root.

The DFS computes subtree heights bottom-up. At some internal node, the sum of the two largest child heights becomes 7, which is the maximum diameter in the tree.

Final answer:

7

Complexity Analysis

Measure Complexity Explanation
Time O(n) Every node is visited exactly once
Space O(h) Recursive call stack stores one frame per tree level

The algorithm performs a single DFS traversal. Each node contributes constant work besides iterating through its children, and the total number of child relationships across the tree is n - 1.

The recursive stack depth depends on the tree height. In the worst case, a skewed tree may have height n, though typical balanced trees use much less stack space.

Test Cases

# Definition for testing
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):
        self.max_diameter = 0

        def dfs(node):
            if not node:
                return 0

            longest = 0
            second_longest = 0

            for child in node.children:
                child_height = dfs(child) + 1

                if child_height > longest:
                    second_longest = longest
                    longest = child_height
                elif child_height > second_longest:
                    second_longest = child_height

            self.max_diameter = max(
                self.max_diameter,
                longest + second_longest
            )

            return longest

        dfs(root)

        return self.max_diameter

sol = Solution()

# Example 1
root1 = Node(1, [
    Node(3, [Node(5), Node(6)]),
    Node(2),
    Node(4)
])
assert sol.diameter(root1) == 3  # longest path across two branches

# Example 2
root2 = Node(1, [
    Node(2, [
        Node(3, [
            Node(5, [Node(6)]),
            Node(4)
        ])
    ])
])
assert sol.diameter(root2) == 4  # skewed deep structure

# Single node
root3 = Node(1)
assert sol.diameter(root3) == 0  # no edges

# Linear chain
root4 = Node(1, [
    Node(2, [
        Node(3, [
            Node(4)
        ])
    ])
])
assert sol.diameter(root4) == 3  # linked-list shaped tree

# Star-shaped tree
root5 = Node(1, [
    Node(2),
    Node(3),
    Node(4),
    Node(5)
])
assert sol.diameter(root5) == 2  # leaf to leaf through root

# Diameter entirely inside subtree
subtree = Node(2, [
    Node(3, [Node(5)]),
    Node(4, [Node(6)])
])

root6 = Node(1, [subtree])

assert sol.diameter(root6) == 4  # longest path does not require root

# Empty child lists
root7 = Node(1, [])
assert sol.diameter(root7) == 0  # explicit empty children

print("All tests passed.")
Test Why
Example 1 Validates standard branching tree
Example 2 Validates deep skewed structure
Single node Ensures diameter 0 is handled
Linear chain Tests linked-list behavior
Star-shaped tree Tests combining two child paths
Diameter inside subtree Ensures path need not pass through root
Empty child lists Confirms empty arrays work correctly

Edge Cases

Single Node Tree

A tree containing only one node has no edges, so its diameter must be 0.

This case can easily cause bugs if an implementation incorrectly counts nodes instead of edges. The DFS correctly returns height 0 for leaf nodes, and the diameter never increases beyond 0.

Highly Skewed Tree

An N-ary tree can degenerate into a linear chain where every node has exactly one child.

In this case, the diameter equals the depth of the tree minus one. Some implementations incorrectly assume that diameters always combine two branches, but here only one branch exists at every node.

The algorithm handles this naturally because second_longest remains 0.

Diameter Does Not Pass Through Root

The longest path may exist entirely inside one subtree.

A naive solution that only considers paths through the root would fail on such inputs.

This implementation updates the diameter at every node during DFS, ensuring all possible subtree diameters are considered.

Nodes With Many Children

Unlike binary trees, an N-ary node may have many children.

A buggy implementation might sort all child heights unnecessarily or incorrectly process only two children.

This solution efficiently tracks only the top two heights while iterating once through the child list.