Codeforces 855C - Helga Hufflepuff's Cup

855C - Helga Hufflepuff's Cup

题意

要求构建一棵树，树上至多可以存在 \(x\) 个权值为 \(k\) 的重要点，且与重要点连边的点的权值必须小于 \(k\)，问有多少种构树方案。

分析

树形DP。

有 \(dp[u][s][cnt]\)，表示以 \(u\) 为根结点的子树，重要点的数目为 \(cnt\) 时的方案数，其中 \(s=0\) 表示 \(u\) 的权值小于 \(k\) ，\(s=1\) 表示 \(u\) 的权值等于 \(k\) ，\(s=2\) 表示 \(u\) 的权值大于 \(k\)。

code

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MOD = 1e9 + 7;
const int MAXN = 1e5 + 10;
vector<int> G[MAXN];
int n, m, k, x;
int dp[MAXN][3][11];
void dfs(int fa, int u) {
    dp[u][0][0] = k - 1; dp[u][1][1] = 1; dp[u][2][0] = m - k;
    for(auto v : G[u]) {
        if(v != fa) {
            dfs(u, v);
            int tmp[3][11] = {0};
            for(int i = 0; i <= x; i++) {
                for(int j = i; j <= x; j++) {
                    tmp[0][j] = (tmp[0][j] + (ll)dp[u][0][j - i] * ((dp[v][0][i] + dp[v][1][i]) % MOD + dp[v][2][i]) % MOD) % MOD;
                    tmp[1][j] = (tmp[1][j] + (ll)dp[u][1][j - i] * dp[v][0][i] % MOD) % MOD;
                    tmp[2][j] = (tmp[2][j] + (ll)dp[u][2][j - i] * (dp[v][0][i] + dp[v][2][i]) % MOD) % MOD;
                }
            }
            memcpy(dp[u], tmp, sizeof tmp);
        }
    }
}
int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    cin >> n >> m;
    for(int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    cin >> k >> x;
    dfs(1, 1);
    int ans = 0;
    for(int i = 0; i < 3; i++) {
        for(int j = 0; j <= x; j++) {
            ans = (ans + dp[1][i][j]) % MOD;
        }
    }
    cout << ans << endl;
    return 0;
}