Tree (树形+换根dp)

链接：https://ac.nowcoder.com/acm/problem/19782
来源：牛客网

修修去年种下了一棵树，现在它已经有n个结点了。
修修非常擅长数数，他很快就数出了包含每个点的连通点集的数量。
澜澜也想知道答案，但他不会数数，于是他把问题交给了你。

题目描述：问对每个点来说，以该点为根，向外延伸，能形成多少棵子树。自己一个点也算一棵子树。

思路：

首先考虑树形dp求一个根的答案。dp1[u]记为以u为根，和他的儿子以及后代能形成的所有子树数。u的所有儿子v的dp1[v]+1的乘积。

void dfs1(int x,int fa)
{
    //先求一个根的答案
    for(int son:to[x])
    {
        if(son==fa) continue;
        dfs1(son,x);
        dp1[x]=(dp1[x]*(dp1[son]+1))%mod;
    }
}

考虑以一个知道的根节点答案换根dp求其他答案。dp2[u]记为以u为根，和整棵树其他点能形成的所有子树数量。（即最终答案）

当从fa的答案换根到他的儿子son时，通过找规律或某些想法可知：dp2[son]=(dp2[fa] / (dp1[son]+1)+1) *dp1[son]。

 

这有一个很坑的地方，就是(dp1[son]+1)求逆元可能得到0(dp1[son]可能为mod-1)，这时相当于除于0，就会出问题。

当逆元inv为0时，看到dp2[fa]实际是由在树形dp的时候求出的dp1[fa]，而dp1[fa]又等于(他所有儿子dp1的值+1)的乘积。

所以dp2[fa] / (dp1[son]+1)又可以变成fa其他儿子的乘积：fa除son外的其他儿子记brother。

(dp1[brother_1]+1) * (dp1[brother_2]+1) * .....他的所有兄弟的值乘积。

void dfs2(int x,int fa)
{
    //换根算答案
    for(int son:to[x])
    {
        if(son==fa) continue;
        //求儿子的dp2值，放dfs上面
        ll inv=mp(dp1[son]+1,mod-2);
        if(inv!=0)
        {
            dp2[son]=(dp2[x]*inv%mod+1)%mod*dp1[son]%mod;
        }
        else
        {
            //逆元为0，不能直接乘，暴力算贡献
            ll tmp=1;
            for(int brother:to[x])
            {
                if(brother==son||brother==fa) continue;
                tmp=tmp*(dp1[brother]+1)%mod;
            }
            dp2[son]=(tmp+1)*dp1[son]%mod;
        }
        dfs2(son,x);   
    }
}

 

代码：

#include<iostream>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<sstream>
#include<vector>
#include<stack>
#include<deque>
#include<cmath>
#include<map>
#include<queue>
#include<bitset>
//#include<hash_map>
#define sd(x) scanf("%d",&x)
#define lsd(x) scanf("%lld",&x)
#define ms(x,y) memset(x,y,sizeof x)
#define fu(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define all(a) a.begin(),a.end()
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
using namespace std;
//using namespace __gnu_cxx;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
const int maxn=1e6+79;
const int mod=1e9+7;
const int INF=1e9+7;
const double pi=acos(-1);
ll mp(ll x,ll n){ll res=1; while(n){if(n&1) res=res*x%mod;x=x*x%mod;n>>=1;} return res;}

ll dp1[maxn],dp2[maxn];
//树形dp+换根dp
vector<int> to[maxn];
void dfs1(int x,int fa)
{
    //先求一个根的答案
    for(int son:to[x])
    {
        if(son==fa) continue;
        dfs1(son,x);
        dp1[x]=(dp1[x]*(dp1[son]+1))%mod;
    }
}
void dfs2(int x,int fa)
{
    //换根算答案
    for(int son:to[x])
    {
        if(son==fa) continue;
        //求儿子的dp2值，放dfs上面
        ll inv=mp(dp1[son]+1,mod-2);
        if(inv!=0)
        {
            dp2[son]=(dp2[x]*inv%mod+1)%mod*dp1[son]%mod;
        }
        else
        {
            //逆元为0，不能直接乘，暴力算贡献
            ll tmp=1;
            for(int brother:to[x])
            {
                if(brother==son||brother==fa) continue;
                tmp=tmp*(dp1[brother]+1)%mod;
            }
            dp2[son]=(tmp+1)*dp1[son]%mod;
        }
        dfs2(son,x);   
    }
}
int main()
{
    int n;sd(n);
    fu(i,1,n-1)
    {
        int u,v;sd(u);sd(v);
        to[u].push_back(v);
        to[v].push_back(u);
    }
    fu(i,1,n) 
    {
        //自己一个点可以算一棵子树
        dp1[i]=1;
    }
    dfs1(1,0);
    dp2[1]=dp1[1];//从1开始换根
    dfs2(1,0);
    fu(i,1,n) printf("%lld\n",dp2[i]);
    return 0;
}