LeetCode-Python-1245. 树的直径

给你这棵「无向树」，请你测算并返回它的「直径」：这棵树上最长简单路径的 边数。

我们用一个由所有「边」组成的数组 edges 来表示一棵无向树，其中 edges[i] = [u, v] 表示节点 u 和 v 之间的双向边。

树上的节点都已经用 {0, 1, ..., edges.length} 中的数做了标记，每个节点上的标记都是独一无二的。

 

示例 1：

输入：edges = [[0,1],[0,2]]
输出：2
解释：
这棵树上最长的路径是 1 - 0 - 2，边数为 2。

示例 2：

输入：edges = [[0,1],[1,2],[2,3],[1,4],[4,5]]
输出：4
解释： 
这棵树上最长的路径是 3 - 2 - 1 - 4 - 5，边数为 4。

 

提示：

• 0 <= edges.length < 10^4
• edges[i][0] != edges[i][1]
• 0 <= edges[i][j] <= edges.length
• edges 会形成一棵无向树

思路：

除了要手动建树， 基本类似 LeetCode-Python-543. 二叉树的直径。

from collections import defaultdict        
class Solution(object):
    def treeDiameter(self, edges):
        """
        :type edges: List[List[int]]
        :rtype: int
        """
        if not edges:
            return 0

        self.neibors = defaultdict(set)
        self.res = 0
        
        for start, end in edges: # 建树
            self.neibors[start].add(end)
        
        def getHeight(node):
            res = []
            for neibor in self.neibors[node]:
                res.append(getHeight(neibor))
            
            while len(res) < 2: # 如果孩子少于两个，就给它补空的上去
                res.append(0)
                
            res = sorted(res)
            self.res = max(self.res, sum(res[-2:])) # 取最长的两个子树长度
            return 1 + max(res)
        
        getHeight(edges[0][0])
        return self.res