牛客 203B tree(树形dp)

大意: 给定树, 对于每个节点, 求包含该节点的连通子集数.

显然有$dp[x]=\prod (dp[y]+1), ans[x]=(\frac{ans[fa[x]]}{dp[x]+1}+1)dp[x]$.

特判$dp[x]+1=0$的情况. 

#include <iostream>
#include <cstdio>
#include <queue>
#define REP(i,a,n) for(int i=a;i<=n;++i)
#define pb push_back
using namespace std;
typedef long long ll;
const int P = 1e9+7;
ll inv(ll x){return x<=1?1:inv(P%x)*(P-P/x)%P;}
  
const int N = 1e6+10;
int n, fa[N];
vector<int> g[N];
int dp[N], dp2[N], ans[N], f[N], cnt[N];
  
void dfs(int x, int f) {
    dp2[x] = dp[x] = 1, fa[x] = f;
    for (int y:g[x]) if (y!=f) {
        dfs(y,x);
        dp[x] = (ll)dp[x]*(1+dp[y])%P;
        if (dp[y]+1==P) ++cnt[x];
        else dp2[x] = (ll)dp2[x]*(1+dp[y])%P;
    }
}
void dfs(int x) {
    if (x==1) ans[x] = dp[x];
    else if (dp[x]+1==P) {
        ll num = f[fa[x]]+1;
        if (cnt[fa[x]]>1) num = 0;
        else num = num*dp2[fa[x]]%P;
        f[x] = num;
        ans[x] = (num+1)*dp[x]%P;
    }
    else {
        f[x] = ans[fa[x]]*inv(dp[x]+1)%P;
        ans[x] = (f[x]+1ll)*dp[x]%P;
    }
    for (int y:g[x]) if (y!=fa[x]) dfs(y);
}
  
int main() {
    scanf("%d", &n);
    REP(i,2,n) {
        int u, v;
        scanf("%d%d", &u, &v);
        g[u].pb(v),g[v].pb(u);
    }
    dfs(1,0),dfs(1);
    REP(i,1,n) printf("%d\n",ans[i]);
}