Program to find the diameter of a n-ary tree in Python

An n-ary tree is a tree where each node can have any number of children. The diameter of an n-ary tree is the longest path between any two nodes in the tree. This path doesn't necessarily have to pass through the root node.

So, if the input is like ?

then the output will be 3.

The diameter consists of the path 27?14?42?56, which has 3 edges (marked in red in the diagram above).

Algorithm

To find the diameter, we need to ?

• For each node, calculate the depths of its subtrees
• The diameter passing through a node is the sum of the two deepest subtree depths
• Track the maximum diameter found across all nodes

Implementation

class Node:
    def __init__(self, value, children=None):
        self.val = value
        self.children = children if children else []

def find_diameter(root):
    diameter = 0
    
    def depth(node):
        nonlocal diameter
        
        if not node:
            return 0
            
        # Get depths of all child subtrees
        child_depths = []
        for child in node.children:
            child_depths.append(depth(child))
        
        # Sort to get the two deepest subtrees
        child_depths.sort(reverse=True)
        
        # Add default 0s if fewer than 2 children
        while len(child_depths) < 2:
            child_depths.append(0)
        
        # Diameter through this node is sum of two deepest subtrees
        current_diameter = child_depths[0] + child_depths[1]
        diameter = max(diameter, current_diameter)
        
        # Return depth of this subtree
        return (child_depths[0] if child_depths else 0) + 1
    
    depth(root)
    return diameter

# Create the n-ary tree
node_65 = Node(65)
node_56 = Node(56)
node_42 = Node(42, [node_56, node_65])
node_32 = Node(32)
node_27 = Node(27)
root = Node(14, [node_27, node_32, node_42])

print("Diameter of the n-ary tree:", find_diameter(root))

Diameter of the n-ary tree: 3

How It Works

The algorithm uses a recursive depth-first approach ?

1. Calculate subtree depths: For each node, recursively find the depth of all child subtrees
2. Find longest path: The diameter through any node is the sum of its two deepest child subtrees
3. Track maximum: Keep track of the maximum diameter found across all nodes
4. Return depth: Each recursive call returns the depth of its subtree for parent calculations

Time Complexity

Conclusion

The diameter of an n-ary tree is found by calculating the sum of the two deepest subtrees at each node. This approach efficiently finds the longest path in O(n) time by visiting each node once.