2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个节点的父节点由一个数组 parent 给出,根节
2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个节点的父节点由一个数组 parent 给出,根节点的父节点为 -1,其他节点的父节点编号一定小于该节点本身。同时,每个节点上还有一个整数值,存放在数组 nums 中,另外给定一个整数 k。
我们需要统计所有满足以下两个条件的非空节点集合的数量:
集合中所有节点的数值之和能够被 k 整除;
集合中不能同时包含任意一个节点和它的直接父节点,也就是说选出的节点在树中不能有相邻的父子关系。
最终结果需要对 1000000007 取模后输出。
n == parent.length == nums.length
1 <= n <= 1000
parent[0] == -1
对于所有的 1 <= i < n:
0 <= parent[i] < i
1 <= nums[i] <= 1000000000
1 <= k <= 100
parent 表示一棵有效的有根树。
输入: parent = [-1,0,0,0], nums = [2,1,2,1], k = 3。
输出: 2。
解释:
有效的子集有:
{1, 2}:节点 1 和 2 都是节点 0 的子节点,且彼此不直接相连。它们的值之和为 1 + 2 = 3 ,可以被 3 整除。
{2, 3}:节点 2 和 3 也不相邻。它们的值之和为 2 + 1 = 3 ,可以被 3 整除。
没有其他子集同时满足两个条件。因此,答案是 2 。
题目来自力扣3939。
合并子节点的详细过程
假设我们已经递归计算了某个子节点 y 的状态fy0和fy1,现在要将 y 合并到当前节点 x 的f0和f1中。
1. 更新f0(不选 x)
此时,由于 x 未被选中,子节点 y可以被选,也可以不被选。因此,从 y 子树中选取的合法集合有两种情况:
- y 不被选,对应
fy0; - y 被选,对应
fy1。
所以,子节点 y 对整体余数的贡献总和为v[i] = fy0[i] + fy1[i](每种余数 i 的方案数相加)。
现在,当前已有的不选 x 的方案数为f0(这是已经处理完之前若干个兄弟子树的累计结果)。当我们把 y 的贡献合并进来时,相当于将两个“余数分布”进行卷积:新余数 = (i + j) % k,其中 i 来自子节点 y 贡献的余数,j 来自之前已处理的子树贡献的余数。新的方案数累加到nf0[(i+j)%k]中。
最后,用nf0替换原有的f0。
2. 更新f1(选 x)
此时,x 已被选中,那么其直接子节点 y 绝对不能选(因为 y 是 x 的子节点,二者相邻)。因此,y 子树只能提供 y 不被选时的方案,即fy0。
类似地,将fy0与当前已有的f1(已经处理完的兄弟子树)进行卷积,得到新的nf1,并替换原有的f1。
递归过程说明
- 整棵树通过
parent数组构建邻接表,根节点为 0。 - 从根节点开始执行深度优先搜索(DFS),递归地处理每个节点。
- 每个节点在处理完所有子节点后,返回自己的
f0和f1给父节点。 - 父节点在得到子节点的返回结果后,按照上述规则合并。
最终答案的计算
当根节点 0 的递归返回后,我们得到了整棵树的f0和f1(分别对应不选根和选根两种全局状态)。
- 合法的非空集合总数 = (不选根时,余数为 0 的方案数) + (选根时,余数为 0 的方案数)。
- 但这两个方案数中都包含了空集(因为初始的
f0[0]=1就代表空集,而选根时不可能包含空集,所以只有f0里有空集),所以最后需要减去空集这一种方案。
即答案 =(f0[0] + f1[0] - 1) mod MOD,最后取模得到正整数结果。
时间复杂度
- 每个节点在合并其每个子节点时,都需要两层循环分别遍历余数 0 到 k-1,因此每次合并的时间开销为
O(k^2)。 - 树中总共有 n 个节点,每条边对应一次合并操作,边的数量为
n-1。 - 因此总时间复杂度为
O(n · k^2)。 - 在本题限制下,
n ≤ 1000,k ≤ 100,故最多约1000 × 10000 = 1e7次基本运算,完全可行。
额外空间复杂度
- 递归深度最坏情况下为
O(n)(例如链状树)。 - 在递归栈的每一层,每个节点会保存若干个长度为 k 的数组(
f0、f1,以及合并时的临时数组),因此递归路径上同时存在的数组总大小约为O(k)乘以递归深度,即O(n · k)。 - 同时,合并过程中产生的临时数组会在函数返回后自动释放,不会累积。
- 因此,额外空间复杂度为
O(n · k),在给定范围内(n=1000, k=100)约为1e5级别,内存充足。
示例验证(以题目输入为例)
parent = [-1,0,0,0],根为 0,子节点为 1、2、3。nums = [2,1,2,1],k=3。- 叶子节点 1、2、3 分别递归返回。
- 根 0 合并子节点后,最终统计余数为 0 的方案数(减去空集)得到答案 2,即
{1,2}和{2,3}两种有效子集,与题意相符。
总结
该算法利用树形 DP 巧妙地处理了“不相邻”和“和整除 k”两个约束,通过分情况(选/不选当前节点)以及卷积合并子节点的方式,在O(n·k²)时间内完成了统计。代码实现清晰,适合本题的数据规模。
Go完整代码如下:
packagemainimport("fmt")funccountValidSubsets(parent[]int,nums[]int,kint)int{constmod=1_000_000_007n:=len(parent)g:=make([][]int,n)fori:=1;i<n;i++{p:=parent[i]g[p]=append(g[p],i)}vardfsfunc(int)([]int,[]int)dfs=func(xint)([]int,[]int){f0:=make([]int,k)// f0[i] 表示不选 x 时,子树 x 的子集点权和模 k 为 i 的方案数f1:=make([]int,k)// f1[i] 表示选 x 时,子树 x 的子集点权和模 k 为 i 的方案数f0[0]=1f1[nums[x]%k]=1for_,y:=rangeg[x]{fy0,fy1:=dfs(y)// 不选 x,那么 y 可选可不选nf0:=make([]int,k)fori:=rangek{// 枚举从子树 y 中选出的点权和模 k 为 iv:=fy0[i]+fy1[i]ifv==0{// 优化continue}forj,w:=rangef0{// 枚举从之前的子树中选出的点权和模 k 为 js:=(i+j)%k nf0[s]=(nf0[s]+v*w)%mod}}// 选 x,那么 y 不能选nf1:=make([]int,k)fori,v:=rangefy0{// 枚举从子树 y 中选出的点权和模 k 为 iifv==0{// 优化continue}forj,w:=rangef1{// 枚举从 x 以及之前的子树中选出的点权和模 k 为 js:=(i+j)%k nf1[s]=(nf1[s]+v*w)%mod}}f0,f1=nf0,nf1}returnf0,f1}f0,f1:=dfs(0)// 恰好被 k 整除即模 k 为 0,注意减去空集的方案数 1return(f0[0]+f1[0]-1+mod)%mod}funcmain(){parent:=[]int{-1,0,0,0}nums:=[]int{2,1,2,1}k:=3result:=countValidSubsets(parent,nums,k)fmt.Println(result)}Python完整代码如下:
# -*-coding:utf-8-*-importsysdefcountValidSubsets(parent,nums,k):MOD=10**9+7n=len(parent)# 构建邻接表g=[[]for_inrange(n)]foriinrange(1,n):p=parent[i]g[p].append(i)sys.setrecursionlimit(max(1000000,n*2+10))defdfs(x):# f0: 不选当前节点 x 时的方案数(按模 k 分类)# f1: 选当前节点 x 时的方案数(按模 k 分类)f0=[0]*k f1=[0]*k f0[0]=1# 空集f1[nums[x]%k]=1# 只含 x 自身的子集forying[x]:fy0,fy1=dfs(y)# 递归处理子节点# ----- 情况1:不选 x,则子节点 y 可选可不选 -----nf0=[0]*kforiinrange(k):v=(fy0[i]+fy1[i])%MODifv==0:continueforjinrange(k):iff0[j]==0:continues=(i+j)%k nf0[s]=(nf0[s]+v*f0[j])%MOD# ----- 情况2:选 x,则子节点 y 不能选 -----nf1=[0]*kforiinrange(k):v=fy0[i]# 子节点只能取不选 y 的方案ifv==0:continueforjinrange(k):iff1[j]==0:continues=(i+j)%k nf1[s]=(nf1[s]+v*f1[j])%MOD f0,f1=nf0,nf1returnf0,f1 f0,f1=dfs(0)# 根节点结果 = 不选根 + 选根,再减去空集(1 种)return(f0[0]+f1[0]-1)%MOD# 示例测试if__name__=="__main__":parent=[-1,0,0,0]nums=[2,1,2,1]k=3print(countValidSubsets(parent,nums,k))C++完整代码如下:
#include<bits/stdc++.h>usingnamespacestd;constlonglongMOD=1000000007LL;pair<vector<longlong>,vector<longlong>>dfs(intx,constvector<vector<int>>&g,constvector<int>&nums,intk){vector<longlong>f0(k,0),f1(k,0);f0[0]=1;// 不选 x 的空集f1[nums[x]%k]=1;// 选 x 的集合(仅包含 x)for(inty:g[x]){auto[fy0,fy1]=dfs(y,g,nums,k);// 不选 x,则子节点 y 可选可不选vector<longlong>nf0(k,0);for(inti=0;i<k;++i){longlongv=(fy0[i]+fy1[i])%MOD;if(v==0)continue;for(intj=0;j<k;++j){if(f0[j]==0)continue;ints=(i+j)%k;nf0[s]=(nf0[s]+v*f0[j])%MOD;}}// 选 x,则子节点 y 不能选vector<longlong>nf1(k,0);for(inti=0;i<k;++i){longlongv=fy0[i];// 只能选 y 中不选 y 的方案if(v==0)continue;for(intj=0;j<k;++j){if(f1[j]==0)continue;ints=(i+j)%k;nf1[s]=(nf1[s]+v*f1[j])%MOD;}}f0=move(nf0);f1=move(nf1);}return{f0,f1};}intcountValidSubsets(constvector<int>&parent,constvector<int>&nums,intk){intn=parent.size();vector<vector<int>>g(n);for(inti=1;i<n;++i){intp=parent[i];g[p].push_back(i);}auto[f0,f1]=dfs(0,g,nums,k);longlongans=(f0[0]+f1[0]-1)%MOD;// 减去空集if(ans<0)ans+=MOD;return(int)ans;}intmain(){vector<int>parent={-1,0,0,0};vector<int>nums={2,1,2,1};intk=3;intresult=countValidSubsets(parent,nums,k);cout<<result<<endl;return0;}