Kimi LeetCode 3373. 连接两棵树后最大目标节点数目 II Rust实现

以下是 LeetCode 3373. 连接两棵树后最大目标节点数目 II 的 Rust 实现。

解题思路

这道题的核心思路是：树是二分图。

- 对每棵树进行二分图染色（0/1），相邻节点颜色不同
- 颜色相同的节点之间距离为偶数，颜色不同的节点之间距离为奇数
- 节点 `u` 是节点 `v` 的 target 当且仅当路径边数为偶数（包括自身距离0）
- 对于 tree1 的节点 `i`，连接 tree2 后：
- tree1 内部能到达的 target 节点数 = 与 `i` 同色的节点数量
- tree2 中可以选择连接一个节点，使得 tree2 贡献的 target 节点数 = `max(tree2中颜色0的数量, tree2中颜色1的数量)`
- 因此 `answer[i] = cnt1[color1[i]] + max(cnt2[0], cnt2[1])`

Rust 代码

```rust
impl Solution {
pub fn max_target_nodes(edges1: Vec<Vec<i32>>, edges2: Vec<Vec<i32>>) -> Vec<i32> {
// 构建邻接表
fn build(edges: &Vec<Vec<i32>>) -> Vec<Vec<i32>> {
let n = edges.len() + 1;
let mut g = vec![vec![]; n];
for e in edges {
let a = e[0] as usize;
let b = e[1] as usize;
g[a].push(b as i32);
g[b].push(a as i32);
}
g
}

// DFS 二分图染色，同时统计每种颜色的节点数量
fn dfs(g: &Vec<Vec<i32>>, a: usize, fa: i32, c: &mut Vec<i32>, d: i32, cnt: &mut Vec<i32>) {
c[a] = d;
cnt[d as usize] += 1;
for &b in &g[a] {
if b != fa {
dfs(g, b as usize, a as i32, c, d ^ 1, cnt);
}
}
}

let g1 = build(&edges1);
let g2 = build(&edges2);
let n = g1.len();
let m = g2.len();

let mut c1 = vec![0; n]; // tree1 各节点的颜色
let mut c2 = vec![0; m]; // tree2 各节点的颜色
let mut cnt1 = vec![0; 2]; // tree1 两种颜色的计数
let mut cnt2 = vec![0; 2]; // tree2 两种颜色的计数

// 先处理 tree2，获取最优连接策略
dfs(&g2, 0, -1, &mut c2, 0, &mut cnt2);
// 再处理 tree1
dfs(&g1, 0, -1, &mut c1, 0, &mut cnt1);

// tree2 能贡献的最大 target 节点数
let t = cnt2[0].max(cnt2[1]);
let mut ans = vec![0; n];
for i in 0..n {
// tree1 中同色的节点数 + tree2 最优贡献
ans[i] = t + cnt1[c1[i] as usize];
}

ans
}
}
```

复杂度分析

- 时间复杂度: O(n + m)，两次 DFS 遍历
- 空间复杂度: O(n + m)，邻接表和辅助数组

代码下载

[LeetCode 3373 Rust 实现](sandbox:///mnt/agents/output/leetcode_3373_rust.rs)