1. 程式人生 > >牛客國慶集訓派對Day3 B Tree(樹形dp + 組合計數)

牛客國慶集訓派對Day3 B Tree(樹形dp + 組合計數)

題意有點繞,其實就是讓你求一個點能被多少個點集包含,同時這些點集內的點要相互連通。

首先,簡單來說,如果只是計算一個有根樹中任意一個點被多少個只包含它以及它的子樹的點的點集包含,那麼直接普通的樹上統計的trick就可以搞定。但是現在問題是,點集的點可以是其子樹之外的點。我們注意到,對於根來說,它的答案都滿足點集裡面的點都在它的子樹裡面,所以說如果能夠把每一個點都當作根來計算一次,那麼就可以簡化問題。

我們令f[x]表示包含x而且點集裡面的點都在x的子樹裡面的點集的個數。那麼f[x]可以通過迭代以下式子得出:

                \large f[x]=f[x]+f[x]*f[son]=f[x]*(f[son]+1)

可以看到,這個式子最後可以化簡成x的兒子對應點集數加一的乘積,具體來說:

                                               \large f[x]=\prod_{son}(f[son]+1)

進行一次dfs我們就可以求出所有的f[x],而根的答案就是f[x]。之後,根據上面說的,我們要想辦法把每一個點都當作根去計算一次。我們考慮再進行一次dfs,這次dfs滿足在搜尋到點x的時候,其父親fa的答案已經算出,也即fa作為根的時候的答案已經算出。現在考慮計算x的答案,也即以x為根的答案。

如果把x看作根,那麼與原本x的子樹相比,x會多一個兒子,也即fa。根據上面的公式,原本的子樹時候的答案是f[x],最後需要乘上一個(f[fa]+1)。但是這裡f[fa]原本是x的父親,這個裡面的答案已經把x以及其子樹考慮過了,所以我們要把這一部分去掉。去掉的方法也很簡單,直接除以當前的(f[x]+1),也即乘以(f[x]+1)的逆元。那麼,總的來說有轉移方程:

                                          \large f[x]=f[x]*(\frac{f[fa]}{1+f[x]}+1)

這樣,用兩個dfs或者說樹形dp,就可以求出最後的答案。

但是呢,本題特別的坑。注意到在第一遍dfs的時候,算出來的f[x]是一堆數字的乘積,而在第二遍dfs的時候,需要乘以一些逆元,也即把一些原本乘過的東西去掉。但是,如果當某個(f[son]+1)對mod取模之後,恰好等於0的話,f[x]就會一直等於0,在第二遍dfs的時候並不能通過乘以逆元的形式把這個乘以0的影響給去掉,這樣就會產生錯誤。為了解決這個問題,我們對於每一個點維護一個flag,表示它的兒子中,有多少個的f[son]+1對mod取模之後等於0。然後在統計對應f[x]的時候不把等於零的兒子給乘進去。如果有兒子是0,刪掉非零兒子的情況對結果無影響;刪掉為零的兒子,且為零的兒子只有一個,那麼在刪掉這個兒子的時候不用乘以逆元,直接用對應的f[x]即可。如果沒有兒子是0,那麼直接乘以逆元即可,同時由於此時相當於把樹旋轉了,所以還要維護flag的變化。具體見程式碼:

#include<bits/stdc++.h>
#define PI 3.1415926535
#define mod 1000000007
#define LL long long
#define pb push_back
#define lb lower_bound
#define ub upper_bound
#define INF 0x3f3f3f3f
#define sf(x) scanf("%d",&x)
#define sc(x,y,z) scanf("%d%d%d",&x,&y,&z)
#define clr(x,n) memset(x,0,sizeof(x[0])*(n+5))
#define file(x) freopen(#x".in","r",stdin),freopen(#x".out","w",stdout)
using namespace std;

const int N = 1000010;

int ls[N<<1],g[N<<1],nxt[N<<1],n,e;
LL f[N]; int flag[N];

inline void upd(LL &x,LL y){x=x*y%mod;}

void dfs1(int x,int fa)
{
    f[x]=1;
    for(int z=ls[x];z;z=nxt[z])
    {
        int y=g[z];
        if (y==fa) continue;
        dfs1(y,x);
        if ((flag[y]?0:f[y])+1==mod) flag[x]++;     //如果是0兒子,那麼統計flag,不計入乘積
            else upd(f[x],(flag[y]?0:f[y])+1);
    }
}

LL qpow(LL x,LL n)
{
    LL res=1;
    while(n)
    {
        if (n&1) res=res*x%mod;
        x=x*x%mod; n>>=1;
    }
    return res;
}

void dfs2(int x,int fa)
{
    if (fa!=0)
        if ((flag[x]?0:f[x])+1==mod)
        {
            if (flag[fa]==1)            //如果有0兒子,且個數恰好是1,且x恰好就是則個0兒子
                if (f[fa]+1==mod) flag[x]++;        //還是要考慮維護flag
                     else upd(f[x],f[fa]+1);
        } else if (!flag[fa])            //如果沒有0兒子,那麼直接用逆元
        {
            LL tmp=f[fa]*qpow(1+(flag[x]?0:f[x]),mod-2)%mod+1;
            if (tmp==mod) flag[x]++; else upd(f[x],tmp);    //維護flag
        }
    for(int z=ls[x];z;z=nxt[z])
    {
        int y=g[z];
        if (y==fa) continue;
        dfs2(y,x);
    }
}

inline void add(int x,int y)
{
    g[++e]=y;
    nxt[e]=ls[x];
    ls[x]=e;
}

int main()
{
    sf(n);
    for(int i=1;i<n;i++)
    {
        int x,y; sf(x); sf(y);
        add(x,y); add(y,x);
    }
    dfs1(1,0);
    dfs2(1,0);
    for(int i=1;i<=n;i++)
        printf("%lld\n",flag[i]?0LL:f[i]);
    return 0;
}