牛客网暑期ACM多校训练营（第三场）G 数论 + BFS

题目链接

---

题意：
给定一颗$n$个节点的树和$k$种颜色，每个节点可以染任意一种颜色，显然一共有$k^n$种染色方案。
定义一个染色方案的鲜艳度为：任意两个相同颜色的节点间距离的最小值。
最后给定一个整数$D$，询问存在多少种方案满足其鲜艳度等于$D$。

---

思路：
定义$F[D]$为鲜艳度大于等于$D$的方案数，则答案可以表示为：$F[D]-F[D+1]$

现在考虑如何求解$F[D]$
对于一个染色方案，其鲜艳度大于等于$D$，等价于不存在任何距离小于$D$的两个点，其颜色相同
故可以考虑使用$BFS$顺序从根开始染色。
当考虑第$i$个节点的染色时，只需要再次以该节点为中心，搜索跟其距离小于$D$且已经被染色过的节点数，由BFS层次遍历的特性，这些节点的颜色一定互不相同，故$k$减去这些节点的数量，即为第$i$个节点可以使用的染色数，若该数量小于等于$0$，则说明不存在满足条件的合法染色方案。

故双重$BFS$可解。

代码：

#include<cstdio>
#include<cstdlib>
#include<cmath>
#include<cstring>
#include<queue>
#include<algorithm>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
 
const ll mod = 1e9 + 7;
const int A = 1e4 + 10;
class P{
public:
    int v,next;
}G[A<<1];
int head[A],tot,N,K,D;
bool vis[A], Get[A];
queue<int> que1;
queue<pii> que2;
 
void add(int u, int v){
    G[tot].v = v;
    G[tot].next = head[u];
    head[u] = tot++;
}
 
int Get_num(int u, int d){
    memset(Get, 0, sizeof(Get));
    int cnt = 0;
    while (!que2.empty()) que2.empty();
    que2.push(make_pair(u, 0));
    Get[u] = 1;
    while (!que2.empty()) {
        pii now = que2.front();que2.pop();
        if(now.second >= d - 1) continue;
 
        for (int i = head[now.first] ; i != -1 ; i = G[i].next) {
            int v = G[i].v;
            if(!Get[v] && vis[v]){
               cnt++;
               Get[v] = 1;
               que2.push(make_pair(v, now.second + 1));
            }
        }
    }
    return cnt;
}
 
ll Solve(int d){
    ll res = 1;
    memset(vis, 0, sizeof(vis));
    while (!que1.empty()) que1.pop();
    que1.push(1);
    while (!que1.empty()) {
        int u = que1.front();que1.pop();
        ll now = (K - Get_num(u, d)) % mod;
        if(now <= 0) return 0;
        res = 1LL * res * now % mod;
        for (int i = head[u]; i != -1; i = G[i].next) {
            int v = G[i].v;
            if (vis[v]) continue;
            que1.push(v);
        }
        vis[u] = 1;
    }
    return res;
}
 
int main(){
    tot = 0;
    memset(head, -1, sizeof(head));
    scanf("%d%d%d", &N, &K, &D);
    for (int i = 1; i < N; i++) {
        int u,v;
        scanf("%d%d", &u, &v);
        add(u,v);add(v,u);
    }
    ll ans = Solve(D) % mod;
    ans = (ans - Solve(D+1) % mod) % mod;
    if (ans < 0) ans += mod;
    printf("%lld\n", ans);
    return 0;
}