Newcoder 141 G.Coloring Tree（组合数学+bfs序）

Description

给出一棵$n$个节点的树，初始状态每点都没有颜色，现在用$K$种颜色给这棵树染色，要求同色的两点在树上的最短距离为$D$，问染色方案数

Input

第一行三个整数$n,K,D$，之后$n-1$行每行输入一条树边

$(1\le K<N\le 5000,1\le D\le N)$

Output

输出染色方案数，结果模$10^9+7$

Sample Input

2 1 1
1 2

Sample Output

1

Solution

令$Solve(d)$为相同颜色最短距离不小于$d$的染色方案数，那么答案即为$Solve(d)-Solve(d+1)$

预处理一遍树上任意两点的最短距离$d(u,v)$，考虑树的$bfs$序，假设$u$点附近与其距离小于$d$的点中有$x$个点颜色已经确定，若$x\ge k$则无解，此时必然冲突，否则$u$点有$k-x$种颜色可以染，之后$u$颜色确定，考虑$bfs$序下一个点即可，注意用$bfs$序作为染色顺序的原因是该种计数需要在深度和层次上都有序，否则会计算错误

Code

#include<cstdio>
#include<algorithm>
#include<vector>
using namespace std;
typedef long long ll;
#define maxn 5005
#define mod 1000000007
int n,K,D,a[maxn];
short d[maxn][maxn];
vector<int>g[maxn];
void dfs(int u,int fa,int w)
{
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i];
        if(v==fa)continue;
        d[v][w]=d[u][w]+1;
        dfs(v,u,w);
    }
}
bool cmp(int x,int y)
{
    return d[1][x]<d[1][y];
}
int Solve(int x)
{
    for(int i=1;i<=n;i++)a[i]=i;
    sort(a+1,a+n+1,cmp);
    int ans=1;
    for(int i=1;i<=n;i++)
    {
        int num=0;
        for(int j=1;j<i;j++)
            if(d[a[i]][a[j]]<x)num++;
        if(num>=K)return 0;
        ans=(ll)ans*(K-num)%mod;
    }
    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);
        g[u].push_back(v),g[v].push_back(u);
    }
    for(int i=1;i<=n;i++)dfs(i,0,i);
    int ans=(Solve(D)-Solve(D+1)+mod)%mod;
    printf("%d\n",ans);
    return 0;
}