CodeForces 855C Helga Hufflepuff‘s Cup 树形dp

传送门：

题意

$给一棵树，有m种颜色1-m，其中x种是特殊颜色，k是特殊颜色的编号。$
$问如何给树涂色，特殊颜色的节点不能连特殊颜色，只能连比他编号小的颜色。$

$问，能有多少种涂色方案？$

思路

$树、节点、计数、特殊情况--树形dp。$

$颜色数量高达1e9，而特殊颜色数量只有10，所以我们围绕特殊颜色展开状态转移。$

$因为特殊颜色的连接是有要求的，所以可以分成三类。$

• $0表示特殊颜色$
• $1表示比特殊颜色编号小的颜色$
• $2表示比特殊颜色编号大的颜色$

$所以设f[u][i][j]表示以u为根节点的子树中，i个特殊颜色，当前点颜色为j的方案数。$

$根据特殊颜色来转移：$
$$f[u][i+j][0]+=f[u][i][0]\ast f[v][j][1]$$

$$f[u][i+j][1]+=f[u][i][1]\ast (f[v][j][0]+f[v][j][1]+f[v][j][2])$$

$$f[u][i+j][2]+=f[u][i][2]\ast (f[v][j][1]+f[v][j][2])$$

$因为转移的时候是+=,所以需要一个辅助函数来帮忙存储，防止新产生的值对后续转移产生影响。$

Code

#include "bits/stdc++.h"
using namespace std;

typedef long long ll;

const ll mod = 1e9 + 7;
const int N = 1e5 + 10;
int n, m, k, x;
vector<int> g[N];
ll dp[11][3], f[N][11][3]; // [以u为根的子树][i种特殊颜色][0/1/2，特殊颜色，小于特殊颜色，大于特殊颜色]
int siz[N];

void dfs(int u, int fa) {
    siz[u] = 1;
    f[u][1][0] = 1;
    f[u][0][1] = k - 1; // 小于特殊颜色的有k-1种
    f[u][0][2] = m - k; // 大于特殊颜色的有m-k种
    for(auto v : g[u]) {
        if(v == fa) continue;
        dfs(v, u);
        mem(dp, 0);
        siz[u] += siz[v];
        for(int i = 0;i <= min(siz[u], x); i++) {
            for(int j = 0;j <= siz[v] && i + j <= x; j++) {
                dp[i + j][0] = (dp[i + j][0] + f[u][i][0] * f[v][j][1] % mod) % mod;
                dp[i + j][1] = (dp[i + j][1] + f[u][i][1] * (f[v][j][0] + f[v][j][1] + f[v][j][2]) % mod) % mod;
                dp[i + j][2] = (dp[i + j][2] + f[u][i][2] * (f[v][j][1] + f[v][j][2]) % mod) % mod;
            }
        }
        for(int i = 0;i <= min(siz[u], x); i++) {
            for(int j = 0;j < 3; j++) {
                f[u][i][j] = dp[i][j];
            }
        }
    }
}

void solve() {
    scanf("%d%d",&n,&m);
    for(int i = 1;i < n; i++) {
        int u, v; scanf("%d%d",&u,&v);
        g[u].eb(v); g[v].eb(u);
    }
    scanf("%d%d",&k,&x);
    dfs(1, 0);
    ll ans = 0;
    for(int i = 0;i <= x; i++) {
        for(int j = 0;j < 3; j++) {
            ans = (ans + f[1][i][j]) % mod;
        }
    }
    printf("%lld\n",ans);
}

signed main() {
    solve();
}