Kimi LeetCode 3544. 子树反转和 Rust实现

📅 2026/7/12 0:15:43 👁️ 阅读次数 📝 编程学习
Kimi    LeetCode 3544. 子树反转和 Rust实现

这是 LeetCode 3544. 子树反转和 的 Rust 实现。

思路

本题是树形 DP + 记忆化搜索。

核心观察: 反转一个节点会影响其整个子树。约束条件要求:若两个反转节点存在祖先关系,则它们之间的距离(边数)必须至少为 `k`。

状态定义: `dp(u, steps, inverted)` 表示:
- 以节点 `u` 为根的子树
- `steps`:距离最近一个祖先反转节点已经经过了多少条边(上限为 `k`)
- `inverted`:当前节点是否被反转(即当前子树整体是否被乘了 `-1`)

转移:
- 若 `steps == k`:当前节点可以再次反转,取 `max(不反转, 反转)`
- 若 `steps < k`:当前节点不能反转,只能继承当前状态
- 对于子节点 `v`:
- 不反转 `u`:子节点继承 `min(steps + 1, k)`,反转状态不变
- 反转 `u`(若允许):子节点的 `steps` 重置为 `1`,反转状态翻转

初始状态: 根节点 `0` 的 `steps = k`,表示根节点可以立即反转。

复杂度: `O(n · k)` 状态数,每个状态遍历所有子节点,总复杂度 `O(n · k)`。

---

```rust
impl Solution {
pub fn subtree_inversion_sum(edges: Vec<Vec<i32>>, nums: Vec<i32>, k: i32) -> i64 {
let n = edges.len() + 1;
let k = k as usize;

// 构建邻接表
let mut graph = vec![vec![]; n];
for edge in edges {
let u = edge[0] as usize;
let v = edge[1] as usize;
graph[u].push(v);
graph[v].push(u);
}

// memo[u][steps][inverted],None 表示未计算
let mut memo = vec![vec![vec![None; 2]; k + 1]; n];
// parent[i] = n 表示无父节点(根节点)
let mut parent = vec![n; n];

fn dfs(
u: usize,
steps: usize,
inverted: usize,
graph: &Vec<Vec<usize>>,
nums: &Vec<i32>,
k: usize,
parent: &mut Vec<usize>,
memo: &mut Vec<Vec<Vec<Option<i64>>>>,
) -> i64 {
if let Some(val) = memo[u][steps][inverted] {
return val;
}

let inv = inverted == 1;
// 当前节点在 inherited 状态下的值
let num = if inv { -(nums[u] as i64) } else { nums[u] as i64 };
// 如果当前节点再反转一次(符号翻转)
let mut neg_num = -num;
let mut sum = num;

for &v in &graph[u] {
if v == parent[u] {
continue;
}
parent[v] = u;

// 不反转当前节点:子节点继承 steps+1(上限 k),反转状态不变
sum += dfs(v, (steps + 1).min(k), inverted, graph, nums, k, parent, memo);

// 如果距离足够远(steps == k),可以尝试反转当前节点
if steps == k {
neg_num += dfs(v, 1, 1 - inverted, graph, nums, k, parent, memo);
}
}

// 如果当前可以反转,取最大值;否则只能不反转
let res = if steps == k { sum.max(neg_num) } else { sum };
memo[u][steps][inverted] = Some(res);
res
}

dfs(0, k, 0, &graph, &nums, k, &mut parent, &mut memo)
}
}
```

---

验证:
- 示例 1:`edges=[[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], nums=[4,-8,-6,3,7,-2,5], k=2` → 27 ✓
- 示例 2:`edges=[[0,1],[1,2],[2,3],[3,4]], nums=[-1,3,-2,4,-5], k=2` → 9 ✓
- 示例 3:`edges=[[0,1],[0,2]], nums=[0,-1,-2], k=3` → 3 ✓