牛客国庆集训派对Day3 - B Tree （树形DP）

https://www.nowcoder.com/acm/contest/203/B

 

题解

 

 

思路差不多，不过没用那个flag。如果碰到dp[u]+1 % mod ==0的时候。直接重新搜一遍。

#include <vector>
#include <stdio.h>
#include <iostream>
#include <stdio.h>
#include <math.h>
using namespace std;
#define LL long long
const LL N = 1e6+6;
const LL mod = 1e9+7;

vector<int>G[N];
LL dp[N],pre[N];
LL ans[N],fdp[N];

void dfs1(LL u,LL p)
{
	pre[u]=p;
	dp[u]=1;
	for(LL i=0;i<G[u].size();i++){
		LL v=G[u][i];
		if(v==p) continue;
		dfs1(v,u);
		dp[u]=dp[u]*(dp[v]+1)%mod;
	}
}

LL qkm(LL a,LL b)
{
	LL ans=1;
	while(b){
		if(b&1) ans=ans*a%mod;
		a=a*a%mod;
		b>>=1;
	}
	return ans;
}


void dfs2(LL u,LL p)
{
	if(p==-1) ans[u]=dp[u];
	else{
		if((dp[u]+1)%mod==0){
			LL cnt=fdp[p]+1;
			for(LL i=0;i<G[p].size();i++){
				LL v=G[p][i];
				if(v==pre[p]) continue;
				if(v==u) continue;
				cnt=cnt*(dp[v]+1)%mod;
			}
			fdp[u]=cnt;
			ans[u]=(cnt+1)*dp[u]%mod;
		}else{
			LL cnt=ans[p]*qkm(dp[u]+1,mod-2)%mod;
			ans[u]=dp[u]*(cnt+1)%mod;
		}
	}
	for(LL i=0;i<G[u].size();i++){
		if(G[u][i]==p) continue;
		dfs2(G[u][i],u);
	}
}


int main()
{
	LL n;
	scanf("%lld",&n);
	for(LL i=1;i<n;i++){
		LL u,v;
		scanf("%lld%lld",&u,&v);
		G[u].push_back(v);G[v].push_back(u);
	}
	dfs1(1,-1);
	dfs2(1,-1);
	for(LL i=1;i<=n;i++){
		printf("%lld\n",ans[i]);
	}
	


}