多叉树的直径

from collections import defaultdict
class Solution:
    def solve(self , n: int, Tree_edge: List[Interval], Edge_value: List[int]) -> int:
        def dfs(node, parent, graph):
            max_distance = 0
            max_node = node
            for neighbor, weight in graph[node]:
                if neighbor == parent:  # 避免回到父节点
                    continue
                distance, cur_max_node = dfs(neighbor, node, graph)
                distance += weight
                if distance > max_distance:
                    max_distance = distance
                    max_node = cur_max_node
            return max_distance, max_node

        graph = defaultdict(list)
        for i in range(len(Tree_edge)):
            u, v, weight = Tree_edge[i].start, Tree_edge[i].end, Edge_value[i]
            graph[u].append((v, weight))
            graph[v].append((u, weight))
        # 第一次 DFS 从任意节点开始，找到最远的节点
        start_node = Tree_edge[0].start
        _, max_node = dfs(start_node, None, graph)
        # 第二次 DFS 从找到的最远节点开始，计算直径
        max_distance, _ = dfs(max_node, None, graph)
        
        return max_distance