Diameter of N-Ary Tree Python Solution

lcpy gen -n 1522   # by problem number
lcpy gen -s diameter_of_n_ary_tree   # by problem name

Input: root = [1,null,3,2,4,null,5,6]
Output: 3
Explanation: Diameter is shown in red color.

Input: root = [1,null,2,null,3,4,null,5,null,6]
Output: 4

Input: root = [1,null,2,3,4,5,null,null,6,7,null,8,null,9,10,null,null,11,null,12,null,13,null,null,14]
Output: 7

from __future__ import annotations


class NaryNode:
    def __init__(self, val: int = 0, children: list[NaryNode] | None = None) -> None:
        self.val = val
        self.children = children if children is not None else []


class Solution:
    # Time: O(n)
    # Space: O(h)
    def diameter(self, root: NaryNode | None) -> int:
        ans = 0

        def dfs(node: NaryNode | None) -> int:
            nonlocal ans
            if node is None:
                return 0
            first = second = 0
            for child in node.children:
                depth = dfs(child)
                if depth > first:
                    second, first = first, depth
                elif depth > second:
                    second = depth
            ans = max(ans, first + second)
            return 1 + first

        dfs(root)
        return ans