[Nowcoder 2018ACM多校第三场G] Coloring Tree

题目大意：
给你一颗节点数为n的树，有K种颜色，每个节点可以染上某一种颜色。定义一种染色方案的染色值为任意两个颜色相同节点距离的最小值。求染色值恰好等于D的染色方案数模 $10^9+7$ 。 $(n, K \leq 5000, D \leq n)$

题目思路：
一个染色值为D的方案，说明每对相同颜色节点的距离大于等于D，同时得至少有一对点取值正好为D，这样不是很好求。
将问题作转化，令F(D)表示染色值大于等于D的方案数，则答案为F(D)-F(D+1)。对于染色值大于等于D的方案，只要满足每对相同颜色节点距离大于等于D即可，取个逆否命题就是没有哪对颜色相同节点距离小于D。
考虑按bfs顺序对节点染色，当枚举到节点u时，考虑所有距离u小于D的已染色点构成的集合，他都不能去选择那些颜色。同时由于我们是按bfs顺序染色的，故所有距离u小于D的已染色点构成的集合中两两点之间的距离也一定是小于D的，故该集合中所有点的颜色两两不同。所以u能染的颜色方案数为K-集合大小。如果小于0直接return 0。否则把每个点的方案数相乘即可。找距离某个节点u小于D的点个数也是可以用bfs或dfs。复杂度 $O(n^2)$

Code:

#include <map>
#include <set>
#include <map>
#include <bitset>
#include <cmath>
#include <queue>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <algorithm>

#define ll long long
#define db double
#define pw(x) ((x) * (x))
#define fi first
#define se second
#define mp(x, y) make_pair(x, y)

using namespace std;

const int N = 1010;
const int mo = (int)1e9 + 7;

int n, D, K;
int cnt, lst[N], nxt[N * 2], to[N * 2];
int head, tail, que[N]; bool col[N];

void add(int u, int v){
    nxt[++ cnt] = lst[u]; lst[u] = cnt; to[cnt] = v;
    nxt[++ cnt] = lst[v]; lst[v] = cnt; to[cnt] = u;
}

int dfs(int u, int fa, int dep, int D){
    if (dep == D) return 0;
    int ret = col[u];
    for (int j = lst[u]; j; j = nxt[j]){
        int v = to[j];
        if (v == fa) continue;
        ret += dfs(v, u, dep + 1, D);
    }
    return ret;
}

ll solve(int D){
    ll ret = 1;
    head = tail = 0;
    que[tail = 1] = 1;
    memset(col, 0, sizeof(col));

    while (head < tail){
        int u = que[++ head];
        int x = dfs(u, 0, 0, D);
        col[u] = 1;
        (ret *= (K - x)) %= mo;
        if (K <= x) return 0;
        for (int j = lst[u]; j; j = nxt[j]){
            int v = to[j];
            if (col[v]) continue;
            que[++ tail] = v;
        }
    }

    return ret;
}

int main(){

    scanf("%d%d%d", &n, &K, &D);
    for (int i = 2, u, v; i <= n; i ++){
        scanf("%d %d", &u, &v);
        add(u, v);
    }

    printf("%lld\n", ((solve(D) - solve(D + 1)) % mo + mo) % mo);

    return 0;
}

[Nowcoder 2018ACM多校第三场G] Coloring Tree

猜你喜欢