问题描述给定一棵 n 个节点的无向树节点编号 0 到 n-1以及一个长度相同的数组 groupgroup[i] 表示节点 i 的分组标签。两个节点 u 和 v 若 group[u] group[v]则它们属于同一组。交互代价定义为树上两节点之间唯一路径的边数。要求返回所有同组无序节点对的交互代价总和。核心思路边贡献统计法直接枚举所有同组节点对并计算路径长度时间复杂度为 O(n²)对于 n ≤ 10⁵ 会超时。核心转化总代价 每条边被同组节点对经过的次数之和。对于任意一条边若将其从树中移除树会被分成两部分。假设某组在这条边的一侧子树中有 x 个节点该组总共有 k 个节点则该组中路径经过这条边的节点对数量为 x * (k - x)。因此只需一次 DFS统计每个子树中各分组的节点数量累加每条边的贡献即可。Java 实现javaimport java.util.ArrayList;import java.util.List;class Solution {private long totalCost 0;private int[][] counts; // counts[u][g] 以u为根的子树中分组g的节点数private int[] totalInGroup; // 全树中各分组的总节点数private ListListInteger adj;public long interactionCosts(int n, int[][] edges, int[] group) {// 1. 构建邻接表adj new ArrayList();for (int i 0; i n; i) {adj.add(new ArrayList());}for (int[] edge : edges) {adj.get(edge[0]).add(edge[1]);adj.get(edge[1]).add(edge[0]);}// 2. 统计各分组总节点数分组标签范围为 1 到 20totalInGroup new int[21];for (int g : group) {totalInGroup[g];}// 3. DFS 统计子树中各分组节点数并累加边的贡献counts new int[n][21];dfs(0, -1, group);return totalCost;}private void dfs(int u, int p, int[] group) {// 当前节点自身属于其分组counts[u][group[u]] 1;for (int v : adj.get(u)) {if (v p) continue;dfs(v, u, group);// 对每个分组计算边 (u, v) 的贡献for (int g 1; g 20; g) {if (totalInGroup[g] 2) continue; // 该组不足2个节点无有效节点对long inSubtree counts[v][g]; // 子树v中分组g的节点数long outsideSubtree totalInGroup[g] - inSubtree; // 子树外同组节点数// 该组中路径经过这条边的节点对数量 inSubtree * outsideSubtreetotalCost inSubtree * outsideSubtree;}// 将子树v的统计结果合并到ufor (int g 1; g 20; g) {counts[u][g] counts[v][g];}}}}代码说明1. 数据结构counts[u][g] 存储以 u 为根的子树中分组 g 的节点数量totalInGroup[g] 存储全树中分组 g 的节点总数。2. DFS 遍历从根节点 0 开始递归遍历。对于每个子节点 v先递归处理 v 的子树得到 counts[v][g]。3. 边贡献计算对于边 (u, v)counts[v][g] 是边下方子树中分组 g 的节点数totalInGroup[g] - counts[v][g] 是边上方同组节点数。二者的乘积就是该组中路径经过这条边的节点对数量。4. 结果合并将子树的统计结果累加到父节点 counts[u][g] 中。复杂度分析· 时间复杂度O(n × G)其中 G 是不同分组的数量本题中 G ≤ 20实际为 O(20n)· 空间复杂度O(n × G) 用于存储 counts 数组
