DeepSeek LeetCode 3786. 树组的交互代价总和 Java实现
问题描述
给定一棵 n 个节点的无向树（节点编号 0 到 n-1），以及一个长度相同的数组 group，group[i] 表示节点 i 的分组标签。两个节点 u 和 v 若 group[u] == group[v]，则它们属于同一组。交互代价定义为树上两节点之间唯一路径的边数。要求返回所有同组无序节点对的交互代价总和。
核心思路：边贡献统计法
直接枚举所有同组节点对并计算路径长度，时间复杂度为 O(n²)，对于 n ≤ 10⁵ 会超时。
核心转化：总代价 = 每条边被同组节点对经过的次数之和。
对于任意一条边，若将其从树中移除，树会被分成两部分。假设某组在这条边的一侧子树中有 x 个节点，该组总共有 k 个节点，则该组中路径经过这条边的节点对数量为 x * (k - x)。
因此只需一次 DFS，统计每个子树中各分组的节点数量，累加每条边的贡献即可。
Java 实现
```java
import 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 List<List<Integer>> 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 到 20）
totalInGroup = 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 * outsideSubtree
totalCost += inSubtree * outsideSubtree;
}
// 将子树v的统计结果合并到u
for (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 数组