2018牛客网暑期ACM多校训练第三场 G-Coloring Tree

题目链接：G-Coloring Tree

 

题意：N个节点的树，可以染色K种

要求相同颜色节点的距离为D，求问方案总数。

 

分析：BFS  

构造函数f(D) 表示相同颜色节点距离至少为D的方案总数，那么最后答案就是f(D)-f(D+1)

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

这些节点任意两个之间的距离一定小于 D ,所以这些节点的颜色是互不相同的。

记录总个数为cnt ，则当前节点可以染的颜色种数为k−cnt 。

答案就累乘一下，时间复杂度 O(n2)

 

代码如下：

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <queue>
#include <vector>
using namespace std;

const int mod = 1000000007;
const int maxn = 5000+10;
int n,k,d;
vector<int>g[maxn];
queue<int>que;
int vis[maxn];
int ta[maxn];
int col;

inline int add(int a, int b) {
   a+=b;
   return a>=mod?a-=mod:a;
}

inline int sub(int a, int b){
    a-=b;
    return a<0?a+=mod:a;
}

inline int mul(long long a, int b){
    a*=b;
    return a>=mod?a%mod:a;
}

int dfs(int u, int fa, int dep){
    if (dep<=0) return 1;
    int ret = 1;
    for (int i=0; i<g[u].size(); ++i) {
        if (!vis[g[u][i]] || g[u][i] == fa) continue;
        ret += dfs(g[u][i],u,dep-1);
    }
    return ret;
}

int f(int D){
     for (int i=1; i<=n; ++i) vis[i] = 0;
     while (!que.empty()) que.pop(); que.push(1);
     int res = 1;

     while (!que.empty()) {
        int u = que.front(); que.pop(); vis[u] = 1;
        int cnt = dfs(u,u,D-1)-1;
        if (cnt >= k) return 0;
        res = mul(res,k-cnt);

        for (int i=0; i<g[u].size(); ++i) {
            if (vis[g[u][i]]) continue;
            que.push(g[u][i]);
        }
     }
     return res;
}

int main(){
    while (scanf("%d%d%d",&n,&k,&d) == 3) {
         for (int i=1; i<=n; ++i) g[i].clear();
         int u,v;
         for (int i = 1; i<n; ++i) {
            scanf("%d%d",&u,&v);
            g[u].push_back(v);
            g[v].push_back(u);
         }

         printf("%d\n",sub(f(d),f(d+1)));
    }
    return 0;
}