牛客国庆集训派对Day3 B(搜索 树形dp)

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/qq_25576697/article/details/82940009

链接：https://www.nowcoder.com/acm/contest/203/B
来源：牛客网
 

时间限制：C/C++ 3秒，其他语言6秒
空间限制：C/C++ 1048576K，其他语言2097152K
64bit IO Format: %lld

题目描述

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

输入描述:

第一行一个整数n (1≤ n ≤ 106)，接下来n-1行每行两个整数ai,bi表示一条边 (1≤ ai,bi≤ n)。

输出描述:

输出n行，每行一个非负整数。第i行表示包含第i个点的连通点集的数量对109+7取模的结果。

示例1

输入

复制

6
1 2
1 3
2 4
4 5
4 6

输出

复制

12
15
7
16
9
9

题意：给一棵树，求出所有节点的所处不同连通块的个数。

思路：可以参考https://blog.csdn.net/u013534123/article/details/82934820写的很详细

本来很简单，做两遍dfs来树形dp一下就可以了，但是第一遍dfs的时候导致的ans[i]=0的情况会影响第二遍dfs的时候的计算，所以需要记录下每个节点的ans[j] + 1 == mod 的子节点j的个数，并且操作的时候，这部分不乘进去，保证所有ans[i]！=0.

#include<iostream>
#include<stdio.h>
#include<string.h>
#include<vector>
using namespace std;
typedef long long ll;
const ll mod = 1e9+7;
const int maxn = 1e6+10;
vector<int>tr[maxn];
ll ans[maxn];
int cnt0[maxn];
ll quick_mod(ll a,ll b)
{
    ll ans  = 1;
    while(b)
    {
        if(b&1)ans = (ans * a)%mod;
        b >>= 1;
        a = (a * a)%mod;
    }
    return ans;
}
void dfs(int cur,int f)
{
    ans[cur] = 1;
    for(int i = 0;i<tr[cur].size();i++)
    {
        if(tr[cur][i]==f)continue;
        dfs(tr[cur][i],cur);
        if((cnt0[tr[cur][i]]?0:ans[tr[cur][i]]) + 1==mod)cnt0[cur]++;
        else ans[cur] = (((cnt0[tr[cur][i]]?0:ans[tr[cur][i]]) + 1)%mod * ans[cur])%mod;
    }
}
inline void upd(ll &x,ll y){x=x*y%mod;}
void dfs2(int cur,int f)
{
    if(f!=-1)
    {
        if(!cnt0[f])  //father node has no son_of_0
        {
            ll tmp = (ans[f] * quick_mod((cnt0[cur]?0:ans[cur])+1,mod-2))%mod;
            tmp = (tmp + 1)%mod;
            if(!tmp)cnt0[cur]++;
            else ans[cur] = (ans[cur] * tmp)%mod;
        }
        else if(cnt0[f]==1&&(cnt0[cur]?0:ans[cur])+1==mod)  //father node has exactly one son_of_0 which is just cur
        {
            ll tmp = (ans[f] + 1)%mod;
            if(!tmp)cnt0[cur]++;
            else ans[cur] = (ans[cur] * tmp)%mod;
        }
        else   //father node has son_of_0 besides cur
        {
            ans[cur] = (ans[cur] * 1)%mod;
        }
    }
    for(int i = 0;i<tr[cur].size();i++)
    {
        if(tr[cur][i] == f)continue;
        dfs2(tr[cur][i],cur);
    }
}
int main()
{
    memset(ans,0,sizeof(ans));
    memset(cnt0,0,sizeof(cnt0));
    int n;
    scanf("%d",&n);
    for(int i = 0;i<=n;i++)tr[i].clear();
    int a,b;
    for(int i  = 0;i<n-1;i++)
    {
        scanf("%d%d",&a,&b);
        tr[a].push_back(b);
        tr[b].push_back(a);
    }
    dfs(1,-1);
    dfs2(1,-1);
    for(int i = 1;i<=n;i++)
    {
        printf("%lld\n",cnt0[i]?0:ans[i]);
    }
}