Nowcoder19782.Tree(树形DP+逆元)

给出一颗n个点的树,求每个点的连通点集的数量。

定义\(f(i)\)表示子树\(i\)内的包括\(i\)的连通点集的数量。

那么每个父节点\(u\)都是它的子节点\(f(j)+1\)相乘而成。

换根时,\(ans[v]=ans[u]/(f(v)+1)+1)*f(v)\)

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e6+100;
const int mod=1e9+7;
typedef long long ll;
ll f[maxn],ans[maxn];
vector<int> g[maxn];
void dfs (int x,int pre) {
	f[x]=1;
	for (int y:g[x]) {
		if (y==pre) continue;
		dfs(y,x);
		f[x]=((f[y]+1)*f[x])%mod; 
	}
}
ll qpow (ll a,ll b) {
	ll ans=1;
	ll base=a;
	while (b) {
		if (b&1) ans=(ans*base)%mod;
		base=base*base%mod;
		b>>=1;
	}
	return ans;
}
void dfs1 (int x,int pre) {
	if (pre) {
		if ((f[x]+1)%mod==0) {
			//0不能求逆元,所以重新计算结果 
			dfs(x,0);
			ans[x]=f[x];
		}
		else {
			ans[x]=((ans[pre]*qpow(f[x]+1,mod-2)%mod)+1)*f[x]%mod;
		}
	}
	else {
		ans[x]=f[x];
	}
	for (int y:g[x]) {
		if (y==pre) continue;
		dfs1(y,x);
	}
}
int n;
int main () {
	scanf("%d",&n);
	for (int i=1;i<n;i++) {
		int x,y;
		scanf("%d%d",&x,&y);
		g[x].push_back(y);
		g[y].push_back(x);
	}
	dfs(1,0);
	dfs1(1,0);
	for (int i=1;i<=n;i++) printf("%lld\n",ans[i]);
}
posted @ 2021-03-31 20:33  zlc0405  阅读(33)  评论(0编辑  收藏  举报