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)。---rustimpl Solution {pub fn subtree_inversion_sum(edges: VecVeci32, nums: Veci32, 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: VecVecusize,nums: Veci32,k: usize,parent: mut Vecusize,memo: mut VecVecVecOptioni64,) - 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;// 不反转当前节点子节点继承 steps1上限 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)}}---验证- 示例 1edges[[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], nums[4,-8,-6,3,7,-2,5], k2 → 27 ✓- 示例 2edges[[0,1],[1,2],[2,3],[3,4]], nums[-1,3,-2,4,-5], k2 → 9 ✓- 示例 3edges[[0,1],[0,2]], nums[0,-1,-2], k3 → 3 ✓