1522. Diameter of N‐Ary Tree - cocoder39/coco39_LC GitHub Wiki

1522. Diameter of N-Ary Tree

be careful with how to determine the max and second max from heap for various heap size

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

        def dfs(node):
            if not node:
                return
            
            min_heap = []
            for child in node.children:
                depth = dfs(child)
                heapq.heappush(min_heap, depth)
                if len(min_heap) > 2:
                    heapq.heappop(min_heap)
            if len(min_heap) == 0:
                max_dep, second_dep = 0, 0
            elif len(min_heap) == 1:
                max_dep, second_dep = min_heap[0], 0
            else: # == 2
                max_dep, second_dep = min_heap[1], min_heap[0]
            self.diameter = max(self.diameter, max_dep + second_dep + 1)
            return max_dep + 1
        
        dfs(root)
        return self.diameter - 1

use 2 variables to update max1 max2

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

        def dfs(node):
            if not node:
                return
            
            max_dep_1, max_dep_2 = 0, 0
            for child in node.children:
                depth = dfs(child)
                if depth > max_dep_1:
                    max_dep_1, max_dep_2 = depth, max_dep_1
                elif depth > max_dep_2:
                    max_dep_2 = depth
            
            self.diameter = max(self.diameter, max_dep_1 + max_dep_2 + 1)
            return max_dep_1 + 1
        
        dfs(root)
        return self.diameter - 1