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 ✓