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;

}