2018牛客网暑假ACM多校训练赛（第三场）G Coloring Tree 计数,bfs

原文链接https://www.cnblogs.com/zhouzhendong/p/NowCoder-2018-Summer-Round3-G.html

题目传送门 - 2018牛客多校赛第三场 G

题意

　　给定一个 $n$ 个节点的树，有 $k$ 种颜色。

　　现在让你给每一个节点都染上一种颜色，总共有 $k^n$ 种方法。

　　现在问，在所有染色方案中，使得相同颜色点对之间的最短距离为 $D$ 的有多少种方案。

　　答案对于 $10^9+7$ 取模。

　　$n,k,d\leq 5000$

题解

　　首先我们来算一下相同颜色点对之间的最短距离不小于 $D$ 的情况。

　　我们考虑 bfs 。

　　对于当前节点，我们找出已经访问过的节点中与当前节点的距离小于 $D$ 的节点。

　　由于我们是 bfs 的，可以证明找出的这些节点任意两个之间的距离一定小于 $D$ 。所以这些节点的颜色互不相同。

　　记这些节点的总个数为 $cnt$ ，则当前节点可以染的颜色种数为 $k-cnt$ 。

　　那么答案就累乘一下就可以了。

　　时间复杂度 $O(n^2)$ 。

　　但是我们要求的是等于 $D$ 的情况。

　　怎么做？

　　减掉大于 $D$ 的情况种数即可。

代码

#include <bits/stdc++.h>
using namespace std;
const int N=5005,mod=1e9+7;
struct Gragh{
	static const int M=N*2;
	int cnt,y[M],nxt[M],fst[N];
	void clear(){
		cnt=0;
		memset(fst,0,sizeof fst);
	}
	void add(int a,int b){
		y[++cnt]=b,nxt[cnt]=fst[a],fst[a]=cnt;
	}
}g;
int n,k,D;
int vis[N],Q[N],head,tail;
int cnt1,cnt2;
void dfs(int x,int pre,int d){
	if (d>D+1)
		return;
	if (d<=D)
		cnt1++;
	cnt2++;
	for (int i=g.fst[x];i;i=g.nxt[i])
		if (g.y[i]!=pre&&vis[g.y[i]])
			dfs(g.y[i],x,d+1);
}
int main(){
	scanf("%d%d%d",&n,&k,&D),D--;
	g.clear();
	for (int i=1,a,b;i<n;i++){
		scanf("%d%d",&a,&b);
		g.add(a,b);
		g.add(b,a);
	}
	memset(vis,0,sizeof vis);
	head=tail=0;
	Q[++tail]=1;
	int ans1=1,ans2=1;
	while (head<tail){
		int x=Q[++head];
		vis[x]=1,cnt1=cnt2=-1;
		dfs(x,0,0);
		cnt1=max(k-cnt1,0),cnt2=max(k-cnt2,0);
		ans1=1LL*ans1*cnt1%mod;
		ans2=1LL*ans2*cnt2%mod;
		for (int i=g.fst[x];i;i=g.nxt[i])
			if (!vis[g.y[i]])
				Q[++tail]=g.y[i];
	}
	printf("%d",(ans1-ans2+mod)%mod);
	return 0;
}