2019牛客CSP-S集训营 #1 B-乃爱与城市拥挤程度(树形DP)

Description

Solution

Code

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1e5+10,M=1e9+7,K=11;
ll f[N][K],dp[N][K],add[N][K],mul[N][K];
int n,k,tot,h[N];
struct Edge{int v,nxt;}e[N<<1];
void Add(int a,int b){e[++tot].v=b,e[tot].nxt=h[a],h[a]=tot;}
ll qp(ll a,ll b)
{
	ll r=1;
	while(b){if(b&1)r=r*a%M;a=a*a%M,b>>=1;}
	return r;
}
void dfs1(int x,int fx)
{
	for(int i=0;i<=k;i++) f[x][i]=dp[x][i]=1;
	for(int i=h[x],v;i;i=e[i].nxt)if((v=e[i].v)!=fx)
	{
		dfs1(v,x);
		for(int j=1;j<=k;j++)
		 	f[x][j]=(f[x][j]+f[v][j-1])%M,
		 	dp[x][j]=dp[x][j]*dp[v][j-1]%M;
	}
	for(int i=1;i<=k;i++) dp[x][i]=dp[x][i]*f[x][i]%M;
}
void dfs2(int x,int fx)
{
	add[x][0]=mul[x][0]=1;
	add[x][1]=mul[x][1]=f[x][1]+f[fx][0];
	for(int i=2;i<=k;i++)
		add[x][i]=(add[x][i]+add[fx][i-1]-f[x][i-2]+f[x][i])%M,
		mul[x][i]=mul[fx][i-1]%M*qp(add[fx][i-1]*dp[x][i-2]%M,M-2)%M
				*dp[x][i]%M*qp(f[x][i],M-2)%M
				*add[x][i]%M
				*(add[fx][i-1]-f[x][i-2]+M)%M;
	for(int i=h[x],v;i;i=e[i].nxt)if((v=e[i].v)!=fx) dfs2(v,x);
}
int main()
{
	scanf("%d%d",&n,&k);
	for(int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),Add(u,v),Add(v,u);
	dfs1(1,0);
	for(int i=1;i<=k;i++) add[1][i]=f[1][i],mul[1][i]=dp[1][i];
	for(int i=h[1];i;i=e[i].nxt) dfs2(e[i].v,1);
	for(int i=1;i<=n;i++) printf("%d%c",add[i][k],i==n?'\n':' ');
	for(int i=1;i<=n;i++) printf("%d%c",mul[i][k],i==n?'\n':' ');
}