牛客多校三-G题-Coloring Tree （bfs+计数）

版权声明：本文为原创文章，转载请标明出处。 https://blog.csdn.net/qq_37171272/article/details/81257064

题目链接

题意：有一颗有n个节点的树，现在你有k种颜色，要给树的每个节点上颜色，然后给出一个D，表示上完 颜色后树的colorness为D，colorness的定义为树上颜色相同的点的最短的距离。求colorness等于D的上色方案数。

思路：

我们定义f（D） 为colorness大于等于D的方案数，那么ans=f（D）-f（D+1）；

对于如何求f（D），通过第一个bfs从根向下遍历树，对于bfs顺序前遇到的每个节点，我们在通过另一个bfs去找到和这个节点相邻的所有距离小于D的并且已经上色的节点数num，因为是通过bfs顺序上色的原因，可以发现num个节点的颜色是各不相同的，如果num>=k，说明当前节点已经没有颜色可以上，ans=0，否则ans*=(k-num);

f(D+1)同理；

AC代码：

/*************************************************************************
	> File Name: nkdx3-G.cpp
	> Author: cyh
	> Mail: 
	> Created Time: 2018年07月27日 星期五 10时05分35秒
 ************************************************************************/

#include<bits/stdc++.h>
using namespace std;
#define LL long long
#define ll tree[rt].l
#define rr tree[rt].r
#define rs rt<<1|1
#define ls rt<<1
const int maxn=5e3+10;
const int inf=1e9+10;
const long long INF=2e18+10;
const int mod=1e9+7;
int n,k,d;
vector<int> vec[maxn];
bool v[maxn],vis[maxn],col[maxn];
int bfs2(int x)
{
    queue<pair<int,int> > que;
    que.push(make_pair(x,0));
    memset(v,false,sizeof(v));
    v[x]=true;
    int sum=0;
    while(que.size()){
        int tmp=que.front().first;
        int s=que.front().second;
        que.pop();
        for(int i=0;i<vec[tmp].size();i++){
            if(!v[vec[tmp][i]]&&vis[vec[tmp][i]]){
                if(s+1<d)que.push(make_pair(vec[tmp][i],s+1)),v[vec[tmp][i]]=true;
                if(s+1<d&&col[vec[tmp][i]]) sum++;
            }
        }
    }
    return sum;
}
LL bfs1()
{
    queue<int> que;
    que.push(1);
    vis[1]=true;
    LL ans=1;
    while(que.size()){
        int tmp=que.front();que.pop(); 
        col[tmp]=true;
        int sum=bfs2(tmp);
        if(sum>=k) return 0;
        ans=ans*(k-sum)%mod;
        for(int i=0;i<vec[tmp].size();i++){
            if(!vis[vec[tmp][i]]){
                que.push(vec[tmp][i]);
                vis[vec[tmp][i]]=true;
            }
        }
    }
    return ans;
}
int main()
{
    scanf("%d%d%d",&n,&k,&d);
    for(int i=1;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        vec[u].push_back(v);
        vec[v].push_back(u);
    }
    LL ans=bfs1();
    d++;
    memset(vis,false,sizeof(vis));
    memset(col,false,sizeof(col));
    ans=(ans-bfs1())%mod+mod;
    printf("%lld\n",ans%mod);
	return 0;
}