leetcode-tree

1245. Tree Diameter.py

class Solution:
    def treeDiameter(self, edges: List[List[int]]) -> int:
        """
        The problem is to calculate diameter of tree, given a graph with tree characteristics. 
        
        Note:
            1. The path with same diameter is not unique.
            2. The diameter of a tree is the longest length of a path between two nodes.

            
        two-pass+ DFS approach
        We find the diameter based on two times dfs.
        1. We choose an arbitrary node a in the tree as starting node and find the farthest node b from a.
        2. Then, we find farthest node c from b. The diameter of the ree is distance between b and c.
        
        T:O(n)
        S:O(n)
        """
        def dfs(cur, prev, l):
            if l > self.ans[0]:
                self.ans[0] = l
                self.ans[1] = cur
            for nei in graph[cur]:
                if nei!= prev:
                    dfs(nei, cur, l+1)
                    
        graph = collections.defaultdict(list)
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)
        # two times dfs
        self.ans = [float("-inf"), -1] # [length, farthest_node_from_starting_node]
        dfs(0, -1, 0)
        furthest_node = self.ans[1]
        
        self.ans = [float("-inf"), -1] # [length, farthest_node_from_starting_node]
        dfs(furthest_node, -1, 0)
        return self.ans[0] # diameter