Codeforces855C Helga Hufflepuff's Cup

Description

给出一颗有 \(n(n\le 10^5)\) 个节点的树。有 \(m\le 10^9\) 种颜色，\(k\) 是特殊颜色。现要染色整棵树，满足

• 特殊颜色点不能与特殊颜色点相邻。
• 与特殊颜色点相邻的点的颜色序号小于 \(k\) 。
• 特殊颜色点数不超过 \(x\) 。

求合法染色方案数。答案对 \(10^9+7\) 取模。

Solution

节点颜色其实只有三种

0. 能与特殊点相邻的点。
1. 特殊颜色点。
2. 除特殊点外不能与特殊点相邻的点。

树上 \(\mathrm{dp}\) 。\(dp[i][j][0/1/2]\) 表示在 \(i\) 的子树内有 \(j\) 个特殊点， \(i\) 的节点类型为 \(0/1/2\) 的方案数。

转移时枚举 \(j\) ，以及某个儿子子树内的特殊点数量 \(k\) 。

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

#define N 100001
#define rep(i, a, b) for (int i = a; i <= b; i++)
#define ll long long

inline int read() {
	int x = 0, flag = 1; char ch = getchar(); for (; !isdigit(ch); ch = getchar()) if (ch == '-') flag = -1;
	for (; isdigit(ch); ch = getchar()) x = (x << 1) + (x << 3) + ch - '0'; return x * flag;
}

const ll P = 1e9 + 7;
int n, m, K, X, head[N], tot;
ll f[N][11][3];
struct edge { int v, next; }e[N << 1];
inline void add(int u, int v) { e[++tot].v = v, e[tot].next = head[u], head[u] = tot; }

void dfs(int u, int fa) {
	f[u][0][0] = K - 1, f[u][1][1] = 1, f[u][0][2] = m - K;
	for (int i = head[u], v; i; i = e[i].next) if ((v = e[i].v) != fa) {
		dfs(v, u);
		ll g[11][3] = { 0 };
		for (int j = X; j >= 0; j--) rep(k, 0, j)
			(g[j][0] += (f[v][k][0] + f[v][k][1] + f[v][k][2]) * f[u][j - k][0]) %= P,
			(g[j][1] += f[v][k][0] * f[u][j - k][1]) %= P,
			(g[j][2] += (f[v][k][0] + f[v][k][2]) * f[u][j - k][2]) %= P;
		memcpy(f[u], g, sizeof g);
	}
}

int main() {
	scanf("%d%d", &n, &m);
	rep(i, 2, n) { int u = read(), v = read(); add(u, v), add(v, u); }
	scanf("%d%d", &K, &X);
	dfs(1, 0);
	ll ans = 0; rep(i, 0, X) rep(j, 0, 2) (ans += f[1][i][j]) %= P;
	cout << ans;
	return 0;
}