E89 换根DP AT_dp_v Subtree

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1e5+10;
ll dp1[N],dp2[N],mod,pre[N],suf[N];
vector<int>G[N];
void dfs1(int x,int xfa)
{ 
	dp1[x]=1; 
    vector<int>son;
	for(int y:G[x])if(y!=xfa){
		dfs1(y,x); 
		dp1[x]=dp1[x]*(dp1[y]+1)%mod;
		son.push_back(y); // 将子节点加入集合，方便之后操作
	}
	ll tmp=1;
	for(int i=0;i<son.size();i++)// 预处理前缀积
	{
		pre[son[i]]=tmp;
		tmp=tmp*(dp1[son[i]]+1)%mod;
	}
	
    tmp=1;
	for(int i=son.size()-1;i>=0;i--)// 预处理后缀积
	{
		suf[son[i]]=tmp;
		tmp=tmp*(dp1[son[i]]+1)%mod;
	}
}
void dfs2(int x,int xfa)
{
	if(xfa==0) dp2[x]=1; // 特判根节点
	else dp2[x]=(dp2[xfa]*(pre[x]*suf[x]%mod)%mod+1)%mod;
	for(int y:G[x])if(y!=xfa)
		dfs2(y,x);
}
int main()
{
	int n;scanf("%d%lld",&n,&mod);
	for(int i=1,x,y;i<n;i++){
		scanf("%d%d",&x,&y);
		G[x].push_back(y); 
		G[y].push_back(x);
	} 
	dfs1(1,0); dfs2(1,0);
	for(int i=1;i<=n;i++)printf("%lld\n",dp1[i]*dp2[i]%mod);
	return 0;
}

换根 DP 做法

这题要用换根 DP 解决，换根 DP 也是一种的树形 DP。

由于题目中说的是无根树，我们将其转化为一个以 1 为根的有根树来处理。

设把 $u$ �染成黑色时，在以 $u$ 为根的子树中，染成黑色的节点与 $u$ 构成一个连通块的方案数为 $dp1_u$；在以 $u$ 为根的子树外，则染成黑色的节点与 $u$ 构成一个连通块的方案数为 $dp2_u$。那么最终答案就是 $dp1_u \times dp2_u$。

显然，$dp1_u$ 的转移方程为:

$$dp1_u = \prod_{v \in son(u)} dp1_v + 1$$

其中，$son(u)$ 表示 $u$ 的所有儿子的集合，$dp1_v + 1$ 表示将节点 $v$ 染成黑色和白色的方案数总和。

$dp2_u$ 的转移方程较为难想，要让 $u$ 与外界连通，只有一个中转点，那就是 $u$ 的父亲 $fa$，$fa$ 既连向了 $u$ 的兄弟、又连向了 $u$ 的祖父、曾祖父、叔叔、堂兄弟等，于是知 $dp2_u$ 的状态转移方程为:

$$dp2_u = \left( dp2_{fa} \times \prod_{v \in brother(u)} dp1_v + \text{1} \right) + 1$$

其中，$brother(u)$ 表示 $u$ 的所有兄弟的集合，$dp_v + 1$ 表示将节点 $v$ 染成黑色和白色的方案数总和。整个式子表示将节点 $fa$ 染成黑色和白色的方案数总和（即是否让 $u$ 与外界连通）。

很明显，求 $n$次 $\prod dp_v + 1$ ($v \in brother(u)$) 的复杂度最坏是 $O(n^2)$ 的，TLE，需要优化。由于要取余，且模数不保证为质数，所以优化是不能用除法，或费马小定理的。

考虑预处理前缀积、后缀积，设 $pre_u$ 表示在 $u$ 左边兄弟的 $dp_v + 1$ 值的积，设 $suf_u$ 表示在 $u$ 右边兄弟的 $dp_v + 1$ 值的积，那么 $\prod dp_v + 1$ ($v \in brother(u)$) 就可以转化为 $pre_u \times suf_u$ 的了。