CF855C Helga Hufflepuff's Cup

· · 题解

思路

计数题，考虑 dp。

考虑每个节点转移需要直到的条件：

  1. 儿子节点所染的颜色与 k 的大小关系。
  2. 已经有多少个点的颜色为 k。

那么我们就根据这个设计状态 f_{u,i,0/1/2} 表示以 u 为根的子树已经有 i 个点染成了 k，且 u 染的颜色大于/小于/等于 k。

如果 u 染的颜色小于 k，那么他的儿子 v 随便染什么颜色都可以。

如果 u 染的颜色等于 k，那么 v 染的颜色必须小于 k。

如果 u 染的颜色大于 k，那么 v 染的颜色可以小于 k，也可以大于 k，但不能等于 k。

于是我们列出转移方程：

f_{u, i+j, 0} = \sum_{i=0}^{\min(sz_u,x)}\sum_{j=0}^{\min(sz_v,x-i)} f_{u,i,0} \times (f_{v, j, 0} + f_{v, j, 1} + f_{v, j, 2}) f_{u, i+j, 1} = \sum_{i=0}^{\min(sz_u,x)}\sum_{j=0}^{\min(sz_v,x-i)} f_{u,i,1} \times f_{v, j, 0} f_{u, i+j, 2} = \sum_{i=0}^{\min(sz_u,x)}\sum_{j=0}^{\min(sz_v,x-i)} f_{u,i,0} \times (f_{v, j, 0} + f_{v, j, 2})

需要注意转移的时候 f 会重复，所以要新开一数组记录原来的 f。

代码

#include <bits/stdc++.h>

using namespace std;

const int N = 1e5 + 10, p = 1e9 + 7;
int n, m, k, maxk, sz[N];
long long f[N][20][3], g[20][3];
vector<int> E[N];

void dfs(int u, int fa)
{
    // 边界情况
    f[u][0][0] = k - 1, f[u][1][1] = 1, f[u][0][2] = m - k, sz[u] = 1;
    for (int v : E[u]) if (v != fa)
    {
        dfs(v, u);
        // 处理重复的问题
        for (int i = 0;i <= min(sz[u], maxk);i++)
            g[i][0] = f[u][i][0], g[i][1] = f[u][i][1], g[i][2] = f[u][i][2], 
            f[u][i][0] = f[u][i][1] = f[u][i][2] = 0;
        // 转移
        for (int i = 0;i <= min(sz[u], maxk);i++)
            for (int j = 0;j <= sz[v] && i + j <= maxk;j++)
                f[u][i+j][0] = (f[u][i+j][0] + g[i][0] * (f[v][j][0] + f[v][j][1] + f[v][j][2]) % p) % p, 
                f[u][i+j][1] = (f[u][i+j][1] + g[i][1] * f[v][j][0] % p) % p, 
                f[u][i+j][2] = (f[u][i+j][2] + g[i][2] * (f[v][j][0] + f[v][j][2]) % p) % p;
        sz[u] += sz[v]; // 注意要先转移在加上大小，否则变为 O(n^2)
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);

    cin >> n >> m;
    for (int i = 1, u, v;i < n;i++)
        cin >> u >> v, E[u].emplace_back(v), E[v].emplace_back(u);
    cin >> k >> maxk;
    dfs(1, 0);
    long long ans = 0;
    for (int i = 0;i <= min(maxk, n);i++)
        ans = (ans + f[1][i][0] + f[1][i][1] + f[1][i][2]) % p; // 统计答案
    cout << ans << "\n";

    return 0;
}