Leetcode 1522 Diameter of N-Ary Tree

Problem statement

https://leetcode.com/problems/diameter-of-n-ary-tree/

Solution

Almost identical to problem 1245. Tree Diameter, but here tree is already given, so it is even simpler.

Complexity

Time complexity is O(n), space is O(h).

Code

class Solution:
    def diameter(self, root):
        def dfs(node):
            maxs = [-1, -1]
            for neib in node.children:
                depth = dfs(neib)
                maxs = sorted(maxs + [depth])[-2:]
            self.res = max(self.res, sum(maxs) + 2)
            return maxs[1] + 1
        
        self.res = 0
        dfs(root)
        return self.res