牛客网暑期ACM多校训练营（第三场）G Coloring Tree

链接：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.

题意：

给你一颗树，要你对他进行染色，色彩数有K种，现在要求必须至少有两个相同颜色的节点的距离等于D，其他的大于D。问方案数。

1 ≤ K < N ≤ 5000
1 ≤ D ≤ N

题解：

我们设F(d)为 所有相同颜色的节点中两两之间都大于等于D的方案数。则答案为F(d)-F(d + 1)。

现在来看怎么求F(d)。

利用BFS在树形图上的性质： 

假设当前已经BFS到了节点a， 则所有与节点a距离 <= k， 并且已经被BFS遍历过的节点中， 任意两点的距离不会超过k。

那么我们BFS一遍树， 对于每个BFS到的节点，我们求一下他可以选的颜色数，乘起来即可。假设与他距离<=k并且已经遍历的节点的个数是X， 那么他可选的颜色数即为 (K-X)。如果X >= K ，则不存在染色方案。这个X可以又通过一遍BFS求。

所以复杂度为N^2。

代码：

#include <bits/stdc++.h>
#define pii pair<int,int>
#define mp make_pair
#define fi first
#define se second
#define pb push_back
#define debug(x) cout<<#x<<" = "<<x<<endl;
typedef long long ll;
using namespace std;
const int MAXN = 5e3 + 5;
const int MOD = 1e9 + 7;

vector<int> G[MAXN];
int vis[MAXN], vis2[MAXN], col[MAXN];
int n, k, d;

int findnum(int st, int d) {
    queue<pii> q;
    q.push(mp(st, 0));
    int cnt = 0;
    vis2[st] = st;
    while(!q.empty()) {
        pii now = q.front(); q.pop();
        if(now.se >= d) return cnt;
        for(int i : G[now.fi]) {
            if(vis2[i] == st || now.se + 1 >= d) continue;
            if(!col[i]) continue;
            vis2[i] = st;
            q.push(mp(i, now.se + 1));
            cnt++;
        }
    }
    return cnt;
}

ll bfs(int d) {
    memset(vis, 0, sizeof(vis));
    memset(vis2, 0, sizeof(vis2));
    memset(col, 0, sizeof(col));
    queue<int> q;
    q.push(1);
    ll ans = 1;
    vis[1] = 1;
    while(!q.empty()) {
        int now = q.front(); q.pop();
        int x = findnum(now, d);
        col[now] = 1;
        if(x >= k) return 0;
        ans = ans * (k - x) % MOD;
        for(int i : G[now]) {
            if(vis[i]) continue;
            vis[i] = 1;
            q.push(i);
        }
    }
    return ans;
}


int main() {
#ifdef LOCAL
    freopen ("input.txt", "r", stdin);
#endif
    cin >> n >> k >> d;
    for(int i = 1; i < n; i++) {
   int x, y;
        scanf("%d %d", &x, &y);
        G[     x].pb(y);
        G[y].pb(x);
    }
    ll ans1 = bfs(d), ans2 = bfs(d + 1);
    printf("%lld\n", (ans1 - ans2 + MOD) % MOD);
    return 0;
}