LC 1522 [M] Diameter of N Ary Tree - ALawliet/algorithms GitHub Wiki

n-ary tree, where we have more than just left and right so only want to check the 2 biggest because the rest gotta be smaller

"""
# 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(root):
            # base case for leaf, first store the maximum depth, second is second maximum depth
            first_max = 0
            second_max = 0

            for neighbor in root.children:
                depth = dfs(neighbor)

                if depth > first_max:
                    second_max = first_max
                    first_max = depth
                    # first_max, second_max = depth, first_max

                elif depth > second_max:
                    second_max = depth

            diameter = first_max + second_max
            self.max_diameter = max(self.max_diameter, diameter)

            return first_max + 1
        
        dfs(root)
        
        return self.max_diameter