LeetCode 3939. 统计有根树中不相邻子集的数目 — Rust 实现题目概述给定一棵有根树n ≤ 1000每个节点有权值 nums[i]。要求统计非空子集的数量满足1. 子集中节点权值之和能被 k 整除k ≤ 1002. 子集中任意两个节点在树中不相邻父子不能同时选结果对 10^9 7 取模。解题思路树上 DP对每个节点 u维护两个长度为 k 的 DP 数组状态 含义dp0[mod] 不选节点 u其子树中选出一些不相邻节点和模 k 为 mod 的方案数dp1[mod] 选节点 u其子树中选出一些不相邻节点和模 k 为 mod 的方案数转移- 当前节点不选子节点可选可不选dp0_new[(ab)%k] dp0[a] * (child_dp0[b] child_dp1[b])- 当前节点选子节点不能选dp1_new[(ab)%k] dp1[a] * child_dp0[b]初始化- dp0[0] 1不选当前节点空集- dp1[nums[u] % k] 1选当前节点答案 (dp0_root[0] dp1_root[0] - 1) % MOD减 1 是排除空集。时间复杂度O(n · k²)空间复杂度O(n · k)。---Rust 代码rustuse std::collections::HashMap;const MOD: i64 1_000_000_007;impl Solution {pub fn count_valid_subsets(parent: Veci32, nums: Veci32, k: i32) - i32 {let n parent.len();let k k as usize;// 构建邻接表子节点列表let mut children: VecVecusize vec![vec![]; n];for i in 1..n {let p parent[i] as usize;children[p].push(i);}// DFS 返回 (dp0, dp1)// dp0[mod]: 不选当前节点子树和模k为mod的方案数// dp1[mod]: 选当前节点子树和模k为mod的方案数fn dfs(u: usize, children: VecVecusize, nums: Veci32, k: usize) - (Veci64, Veci64) {let mut dp0 vec![0i64; k];let mut dp1 vec![0i64; k];// 初始化dp0[0] 1; // 不选u空集dp1[(nums[u] as usize) % k] 1; // 选ufor v in children[u] {let (child_dp0, child_dp1) dfs(v, children, nums, k);let mut new_dp0 vec![0i64; k];let mut new_dp1 vec![0i64; k];// 当前节点不选子节点可选可不选for i in 0..k {if dp0[i] 0 { continue; }for j in 0..k {if child_dp0[j] 0 child_dp1[j] 0 { continue; }let ways (child_dp0[j] child_dp1[j]) % MOD;let ni (i j) % k;new_dp0[ni] (new_dp0[ni] dp0[i] * ways) % MOD;}}// 当前节点选子节点不能选for i in 0..k {if dp1[i] 0 { continue; }for j in 0..k {if child_dp0[j] 0 { continue; }let ni (i j) % k;new_dp1[ni] (new_dp1[ni] dp1[i] * child_dp0[j]) % MOD;}}dp0 new_dp0;dp1 new_dp1;}(dp0, dp1)}let (dp0_root, dp1_root) dfs(0, children, nums, k);let ans (dp0_root[0] dp1_root[0] - 1 MOD) % MOD;ans as i32}}---代码要点说明1. 取模处理Rust 中负数取模需要小心最后答案用 (ans MOD) % MOD 保证非负。2. 递归 DFS由于 n ≤ 1000递归深度安全。3. 状态合并对每个子节点做背包式合并复杂度 O(k²)。4. 空集排除dp0[0] 1 表示不选任何节点的空集最终答案需要减 1。
