Hard
题目描述
给你一个整数 n 和一棵有 n 个节点的无向树，节点编号从 0 到 n - 1。树用一个长度为 n - 1 的二维数组 edges 表示，其中 edges[i] = [ui, vi] 表示节点 ui 和 vi 之间有一条无向边。
同时给你一个长度为 n 的整数数组 group，其中 group[i] 表示分配给节点 i 的组标签。
- 如果 group[u] == group[v]，则认为两个节点 u 和 v 属于同一组。
- 节点 u 和 v 之间的交互成本定义为连接它们的树中唯一路径上的边数。
返回所有满足 u != v 且 group[u] == group[v] 的无序对 (u, v) 的交互成本之和。
示例 1：
输入：n = 3, edges = [[0,1],[1,2]], group = [1,1,1]
输出：4
解释：所有节点都属于组 1。节点对之间的交互成本为：
- 节点 (0, 1)：1
- 节点 (1, 2)：1
- 节点 (0, 2)：2
因此，总交互成本为 1 + 1 + 2 = 4。
示例 2：
输入：n = 3, edges = [[0,1],[1,2]], group = [3,2,3]
输出：2
解释：节点 0 和 2 属于组 3，它们之间的交互成本为 2。
约束条件：
- 1 <= n <= 10^5
- edges.length == n - 1
- 0 <= ui, vi <= n - 1
- group.length == n
- 1 <= group[i] <= 20
- 输入保证 edges 表示一棵有效的树
解题思路
这是一道关于树上路径统计的题目。我们需要计算所有相同组别节点对之间的距离总和。
核心思路：
对于树上的每条边，我们可以分析它对答案的贡献。当我们删除一条边时，树被分成两个连通分量。对于同一组别的节点，如果它们分别在两个连通分量中，那么它们之间的路径必须经过这条边，因此这条边对答案的贡献就是两个分量中该组别节点数的乘积。
具体算法：
- 使用后序DFS遍历树，对每个子树统计各组别的节点数量
- 对于每条边 (u, v)，假设 v 是 u 的子节点，计算删除这条边后：
- v 子树中每个组别的节点数：subtree_count[group]
- 剩余部分中每个组别的节点数：total_count[group] - subtree_count[group]
- 该边对组别 g 的贡献为：subtree_count[g] × (total_count[g] - subtree_count[g])
- 将所有边、所有组别的贡献累加即为答案
这种方法的时间复杂度为 O(n)，空间复杂度为 O(n)，是最优解法。
代码实现
class Solution {
public:
long long interactionCosts(int n, vector<vector<int>>& edges, vector<int>& group) {
vector<vector<int>> graph(n);
for (auto& edge : edges) {
graph[edge[0]].push_back(edge[1]);
graph[edge[1]].push_back(edge[0]);
}
// 统计每个组的总节点数
unordered_map<int, int> totalCount;
for (int g : group) {
totalCount[g]++;
}
long long result = 0;
function<unordered_map<int, int>(int, int)> dfs = [&](int node, int parent) -> unordered_map<int, int> {
unordered_map<int, int> subtreeCount;
subtreeCount[group[node]] = 1;
for (int child : graph[node]) {
if (child != parent) {
auto childCount = dfs(child, node);
// 计算边 (node, child) 的贡献
for (auto& [g, count] : childCount) {
long long contribution = (long long)count * (totalCount[g] - count);
result += contribution;
// 合并子树统计
subtreeCount[g] += count;
}
}
}
return subtreeCount;
};
dfs(0, -1);
return result;
}
};
class Solution:
def interactionCosts(self, n: int, edges: List[List[int]], group: List[int]) -> int:
from collections import defaultdict, Counter
# 构建邻接表
graph = defaultdict(list)
for u, v in edges:
graph[u].append(v)
graph[v].append(u)
# 统计每个组的总节点数
total_count = Counter(group)
result = 0
def dfs(node, parent):
nonlocal result
# 当前子树中各组的节点数
subtree_count = defaultdict(int)
subtree_count[group[node]] = 1
for child in graph[node]:
if child != parent:
child_count = dfs(child, node)
# 计算边 (node, child) 的贡献
for g, count in child_count.items():
contribution = count * (total_count[g] - count)
result += contribution
# 合并子树统计
subtree_count[g] += count
return subtree_count
dfs(0, -1)
return result
public class Solution {
public long InteractionCosts(int n, int[][] edges, int[] group) {
var graph = new List<int>[n];
for (int i = 0; i < n; i++) {
graph[i] = new List<int>();
}
foreach (var edge in edges) {
graph[edge[0]].Add(edge[1]);
graph[edge[1]].Add(edge[0]);
}
// 统计每个组的总节点数
var totalCount = new Dictionary<int, int>();
foreach (int g in group) {
totalCount[g] = totalCount.GetValueOrDefault(g, 0) + 1;
}
long result = 0;
Dictionary<int, int> Dfs(int node, int parent) {
var subtreeCount = new Dictionary<int, int>();
subtreeCount[group[node]] = 1;
foreach (int child in graph[node]) {
if (child != parent) {
var childCount = Dfs(child, node);
// 计算边 (node, child) 的贡献
foreach (var kvp in childCount) {
int g = kvp.Key;
int count = kvp.Value;
long contribution = (long)count * (totalCount[g] - count);
result += contribution;
// 合并子树统计
subtreeCount[g] = subtreeCount.GetValueOrDefault(g, 0) + count;
}
}
}
return subtreeCount;
}
Dfs(0, -1);
return result;
}
}
var interactionCosts = function(n, edges, group) {
// 构建邻接表
const graph = Array(n).fill(null).map(() => []);
for (const [u, v] of edges) {
graph[u].push(v);
graph[v].push(u);
}
// 统计每个组的总节点数
const totalCount = new Map();
for (const g of group) {
totalCount.set(g, (totalCount.get(g) || 0) + 1);
}
let result = 0;
function dfs(node, parent) {
const subtreeCount = new Map();
subtreeCount.set(group[node], 1);
for (const child of graph[node]) {
if (child !== parent) {
const childCount = dfs(child, node);
// 计算边 (node, child) 的贡献
for (const [g, count] of childCount) {
const contribution = count * (totalCount.get(g) - count);
result += contribution;
// 合并子树统计
subtreeCount.set(g, (subtreeCount.get(g) || 0) + count);
}
}
}
return subtreeCount;
}
dfs(0, -1);
return result;
};
复杂度分析
| 复杂度类型 | 分析 |
|---|---|
| 时间复杂度 | O(n) - 每个节点访问一次，每条边处理一次 |
| 空间复杂度 | O(n) - 递归栈深度最大为 O(n)，邻接表存储需要 O(n) 空间 |