牛客多校第三场 G Coloring Tree（计数+bfs）

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

题目描述

Christmas is coming! Eddy has received a Christmas tree as gift. Not surprisingly, the tree consists of N vertices and N-1 edges and magically remains connected. Currently, all the vertex of the tree is uncolored. Eddy wants to color each vertex into one of K colors. However, there are too many way to color the tree(i.e. KN ways). Eddy doesn't want the result of coloring being too boring. Thus, he defines the colorness of a tree as follow:

The colorness of a tree is the minimum distance between two vertex colored in the same color.

Now, Eddy is wondering how many way to color the tree such that the colorness of the tree will be D.

输入描述:

The first line of input contains three space-separated integer N, K, D indicating the number of vertices, number of colors, and the required colorness.
For each following N-1 lines, each contains two space-separated positive integer ui, vi indicating that there's an edge between ui and vi.

1 ≤ K < N ≤ 5000
1 ≤ D ≤ N
1 ≤ ui < vi ≤ N
It's guaranteed that the given input is a tree.

输出描述:

Output one line contains an integer indicating the number of way module 1000000007(109+7) to color the tree resulting in colorness being D.

示例1

输入

复制

2 1 1
1 2

输出

复制

1

示例2

输入

复制

4 3 2
1 2
2 3
3 4

输出

复制

18

示例3

输入

复制

4 3 2
1 2
1 3
1 4

输出

复制

24

题目大意：有一棵n个结点的树，现在要对这n个结点进行染色，总共有k中颜色可以选择。现在定义一棵树的颜色度数为两个相同颜色的结点之间的最小距离，现在要你求出颜色度数为d 的染色方案数有多少种，最终结果对1e9+7取模。

题目思路：对于本题，我们定义F(x)为树上颜色相同的两点的最小距离大于x的方案数，那么我们所要求的结果就是F(d-1)-F(d)。

我们要怎么求F(x)呢。我们以bfs对每个点的情况进行遍历，对于每个结点，我们可以做一遍dfs，统计一下与该结点距离小于等于x同时已经染上颜色的结点个数num，如果num>=k，说明无法对当前节点进行染色，就直接返回0，如果num

具体实现看代码：

#include 
#define fi first
#define se second
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define lowbit(x) x&-x
#define MP make_pair
#define pb push_back
#define debug(x) cout<<"x= "< vi;
typedef pairpii;
const int mod=1e9+7;
const ll infll=0x3f3f3f3f3f3f3f3f;
const int MX = 1e5 + 7;
const double EPS=1e-8;
const int INF = 0x3f3f3f3f;
ll qpow(ll a,ll b){
	ll res=1;
	while(b){
		if(b&1) res=(res*a)%mod;
		a=(a*a)%mod;
		b>>=1;
	}
	return res;
}

int n,k,d;
vectorE[MX];
bool vis[MX];
int dfs(int u,int fa,int dep){
	if(!vis[u]) return 0;
	if(dep==0) return 1;
	int res=1;
	for(auto v:E[u]){
		if(v==fa) continue;
		res+=dfs(v,u,dep-1);
	}
	return res;
}
ll solve(int x){
	if(x==0) return qpow(k,n);
	memset(vis,0,sizeof(vis));
	queueq;
	q.push(1);
	ll ans=1;
	while(!q.empty()){
		int u=q.front();q.pop();
		vis[u]=1;
		int cnt=dfs(u,u,x)-1;
		if(cnt>=k) return 0;
		ans=ans*(k-cnt)%mod;
		for(auto v:E[u]){
			if(vis[v]) continue;
			q.push(v);
		}
	}
	return ans;
}

int main(){
	scanf("%d%d%d",&n,&k,&d);
	for(int i=1;i