原文链接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;
}