问题描述给定一棵 n 个节点的无向树节点编号 0 到 n-1数组 group[i] 表示节点 i 所属的分组。两个节点 u 和 v 的交互代价为树上它们之间唯一路径的边数。要求返回所有同组无序节点对的交互代价总和。核心思路边贡献统计直接枚举所有同组节点对并计算路径长度会达到 O(n²)不可行。关键转化总代价等于每条边被同组节点对的路径经过的次数之和。对于任意一条边将其从树中删除会把树分成两部分。假设某一分组在这条边一侧子树中有 x 个节点该组全局总数为 k那么该组中路径经过这条边的节点对数量为 x * (k - x)。因此只需一次 DFS 遍历统计每个子树中各分组的节点数然后累加每条边的贡献即可。Python3 实现pythonfrom typing import Listclass Solution:def interactionCosts(self, n: int, edges: List[List[int]], group: List[int]) - int:# 1. 构建邻接表adj [[] for _ in range(n)]for u, v in edges:adj[u].append(v)adj[v].append(u)# 2. 离散化分组标签若分组标签本身连续且范围小可直接用数组unique_groups sorted(set(group))g_to_idx {g: i for i, g in enumerate(unique_groups)}m len(unique_groups) # 不同分组的数量group_idx [g_to_idx[g] for g in group] # 转为紧凑索引# 3. 统计全局每个分组的节点总数total [0] * mfor gi in group_idx:total[gi] 1# 4. counts[u][gi] 以 u 为根的子树中分组 gi 的节点数counts [[0] * m for _ in range(n)]self.ans 0import syssys.setrecursionlimit(1000000) # 防止递归深度过大def dfs(u: int, p: int) - None:counts[u][group_idx[u]] 1 # 当前节点自身for v in adj[u]:if v p:continuedfs(v, u)# 计算边 (u, v) 对答案的贡献for gi in range(m):if total[gi] 2: # 该组至少两个节点才可能产生贡献c counts[v][gi] # 子树 v 中该组节点数if c 0:self.ans c * (total[gi] - c)# 合并子树的统计信息到 ufor gi in range(m):counts[u][gi] counts[v][gi]dfs(0, -1)return self.ans代码详解1. 离散化处理由于分组标签可能不连续或很大先用 set 收集所有标签并映射为 0..m-1 的整数使后续数组操作高效且紧凑。2. 全局统计total[gi] 记录整个树中分组 gi 的节点总数。3. DFS 遍历从根节点 0 开始递归。对于节点 u· 初始化 counts[u][group_idx[u]] 1。· 遍历每个子节点 v先递归处理 v。· 对每一个分组 gi子树 v 中该组的节点数为 counts[v][gi]则该组中路径经过边 (u, v) 的节点对数为 counts[v][gi] * (total[gi] - counts[v][gi])累加到答案。· 将 v 的计数合并到 u 的计数中。4. 复杂度分析· 时间复杂度O(n · m)其中 m 为不同分组的数量通常很小如题目中 ≤ 20。· 空间复杂度O(n · m) 用于存储子树计数数组加上 O(n) 的邻接表。5. 注意事项· 递归深度可能达到 n已设置 sys.setrecursionlimit 以避免栈溢出。· 若分组标签范围很小例如 1~20可省去离散化步骤直接使用固定大小的数组代码更简洁。测试示例供参考python# 示例# n 4, edges [[0,1],[0,2],[2,3]], group [1,2,1,2]# 同组节点对组1: (0,2) 路径长度 1组2: (1,3) 路径长度 2 → 总和 3sol Solution()print(sol.interactionCosts(4, [[0,1],[0,2],[2,3]], [1,2,1,2])) # 输出 3本解法充分利用了“边贡献”思想将路径总长转化为对每条边的独立贡献并通过一次 DFS 高效完成统计。