带权树的直径（Hard）

• 题目链接

# class Interval:
#     def __init__(self, a=0, b=0):
#         self.start = a
#         self.end = b

#
# 树的直径
# @param n int整型 树的节点个数
# @param Tree_edge Interval类一维数组 树的边
# @param Edge_value int整型一维数组 边的权值
# @return int整型
#
class Solution:
    # 树形dp？
    def __init__(self):
        self.max_distance = float("-inf")

    def solve(self, n, tree_edge, edge_weight):
        from collections import defaultdict
        tree = defaultdict(list)
        # 首先，根据父子关系及边权重构建无向图（无环！）
        for edge, weight in zip(tree_edge, edge_weight):
            tree[edge.start].append((edge.end, weight)) # 存入边尾及权重
            tree[edge.end].append((edge.start, weight)) # 存入边头及权重
            
        self.dfs(tree, 0, defaultdict(bool))
        
        return self.max_distance

    def dfs(self, tree, position, visit):
        left, right = 0, 0
        visit[position] = True # 当前位置设置为已访问
        for child, weight in tree[position]:
            if visit[child]: continue # 若当前位置已经访问过，则跳过
                
            weight += self.dfs(tree, child, visit)
            if weight > left:
                right = left   # 先用right暂存left
                left = weight  # 更新left为更大值
            elif weight > right:
                right = weight # 更新right为更大值
                
            self.max_distance = max(self.max_distance, left + right) # left+right左右路径之和
            
        return max(left, right)