DeepSeek LeetCode 3786. 树组的交互代价总和 C++实现

问题描述给定一棵 n 个节点的无向树节点编号 0 到 n-1数组 group[i] 表示节点 i 所属的分组。两个节点 u 和 v 的交互代价为树上它们之间唯一路径的边数。要求返回所有同组无序节点对的交互代价总和。---核心思路边贡献统计法直接枚举所有同组节点对并计算路径长度会达到 O(n²)不可行。关键转化总代价 每条边被同组节点对的路径经过的次数之和。对于任意一条边将其从树中删除会把树分成两部分。假设某一分组在这条边一侧子树中有 x 个节点该组全局总数为 k则该组中路径经过这条边的节点对数量为 x * (k - x)。因此只需一次 DFS 遍历统计每个子树中各分组的节点数然后累加每条边的贡献即可。---C 实现优化版本cpp#include vector#include functionalusing namespace std;class Solution {public:long long interactionCosts(int n, vectorvectorint edges, vectorint group) {// 1. 构建邻接表vectorvectorint adj(n);for (const auto e : edges) {adj[e[0]].push_back(e[1]);adj[e[1]].push_back(e[0]);}// 2. 分组标签范围已知为 1~20直接使用固定大小数组const int MAX_GROUP 20;vectorint total(MAX_GROUP 1, 0); // 全局各分组节点总数// 3. 统计每个分组的总节点数for (int g : group) {total[g];}// 4. cnt[u][g] 以 u 为根的子树中分组 g 的节点数vectorvectorint cnt(n, vectorint(MAX_GROUP 1, 0));long long ans 0;// DFS 递归函数functionvoid(int, int) dfs [](int u, int parent) {cnt[u][group[u]] 1; // 当前节点自身for (int v : adj[u]) {if (v parent) continue;dfs(v, u);// 计算边 (u, v) 对答案的贡献for (int g 1; g MAX_GROUP; g) {if (total[g] 2) continue; // 该组少于2个节点无贡献long long inSubtree cnt[v][g]; // 子树中该组节点数long long outSubtree total[g] - inSubtree; // 子树外该组节点数if (inSubtree 0 outSubtree 0) {ans inSubtree * outSubtree;}}// 合并子树的统计信息到当前节点for (int g 1; g MAX_GROUP; g) {cnt[u][g] cnt[v][g];}}};dfs(0, -1);return ans;}};---迭代版 DFS避免递归栈溢出cpp#include vector#include stackusing namespace std;class Solution {public:long long interactionCosts(int n, vectorvectorint edges, vectorint group) {// 1. 构建邻接表vectorvectorint adj(n);for (const auto e : edges) {adj[e[0]].push_back(e[1]);adj[e[1]].push_back(e[0]);}// 2. 分组标签范围已知为 1~20const int MAX_GROUP 20;vectorint total(MAX_GROUP 1, 0);for (int g : group) total[g];// 3. cnt[u][g] 统计子树中各分组节点数vectorvectorint cnt(n, vectorint(MAX_GROUP 1, 0));long long ans 0;// 4. 迭代 DFS先序遍历获取父节点关系vectorint parent(n, -1);vectorint order; // 存储遍历顺序stackint st;st.push(0);parent[0] -2; // 标记根节点已访问while (!st.empty()) {int u st.top();st.pop();order.push_back(u);for (int v : adj[u]) {if (v parent[u]) continue;parent[v] u;st.push(v);}}// 5. 逆序遍历从叶子到根累计贡献for (int i n - 1; i 0; --i) {int u order[i];cnt[u][group[u]] 1; // 当前节点自身for (int v : adj[u]) {if (v parent[u]) continue; // 只处理子节点// 已经处理完子节点 v累加它的贡献for (int g 1; g MAX_GROUP; g) {if (total[g] 2) continue;long long inSubtree cnt[v][g];long long outSubtree total[g] - inSubtree;if (inSubtree 0 outSubtree 0) {ans inSubtree * outSubtree;}}// 合并子节点计数for (int g 1; g MAX_GROUP; g) {cnt[u][g] cnt[v][g];}}}return ans;}};---代码说明1. 分组范围优化题目中分组标签范围已知为 1~20直接使用固定大小的数组无需离散化提升效率。2. 两种 DFS 实现· 递归版代码简洁直观适合 n ≤ 10⁵ 的情况需设置栈大小。· 迭代版避免递归栈溢出风险适合大规模数据n ≤ 2×10⁵。3. 核心计算对于边 (u, v)cnt[v][g] 为子树中该组节点数total[g] - cnt[v][g] 为子树外同组节点数。乘积即为该组中路径经过这条边的节点对数量。4. 复杂度分析· 时间复杂度O(n × 20) ≈ O(20n)· 空间复杂度O(n × 20) 用于存储 cnt 数组---测试示例cppint main() {Solution sol;int n 4;vectorvectorint edges {{0,1}, {0,2}, {2,3}};vectorint group {1, 2, 1, 2};cout sol.interactionCosts(n, edges, group) endl; // 输出: 3return 0;}解释· 同组节点对组1: (0,2) 路径长度 1组2: (1,3) 路径长度 2· 总代价 1 2 3---注意事项· 答案可能很大使用 long long 存储结果。· 若分组标签范围不是 1~20需先离散化处理。· 迭代版中order 数组存储了 DFS 前序序列逆序遍历即可保证子节点先于父节点处理。