[数据结构与算法]月之暗面(树形dp)

月之暗面

[Link](D-月之暗面_牛客练习赛97 (nowcoder.com))

题意

给一颗$n$个结点的数染色，普通染色$x$种，特殊染色$y$种，普通颜色无限制，两个相邻结点不能染相同的特殊颜色

求染色方案数

思路

先按当前结点普通颜色还是特殊颜色来分类，记$f[u,0]:以u为根且u填普通颜色的方案数，f[u,1]:以u为根且u填特殊颜色的方案数$，互相独立可以相加，关键在于怎么不从不漏的划分方案。

考虑维护方案对于一个点$u$与它相邻的非父亲结点$v_i$是相互独立的互不干扰符合乘法原理。

分开来看对于$u$选$0$子树无限制转移为$f[u,0]?=f[v_i,1]+f[v_i,0]$，对于$u$选$1$的子树当前根会限制这个子树选$1$的情况，$f[v_i][1]$中$v_i$选$1$一共$y$种不同颜色，且无本质差别，由于$u$的存在$v_i$只能选$y?1$种颜色，因此转移为$f[u,1]?=f[v_i,0]+f[v_i,1]\times \frac{y?1}{y}$

Code

#include <bits/stdc++.h>
#define x first
#define y second
#define debug(x) cout<<#x<<":"<<x<<endl;
using namespace std;
typedef long double ld;
typedef long long LL;
typedef pair<int, int> PII;
typedef pair<double, double> PDD;
typedef unsigned long long ULL;
const int N = 1e6 + 10, M = 2 * N, INF = 0x3f3f3f3f, mod = 998244353;
const double eps = 1e-8, pi = acos(-1), inf = 1e20;
int dx[] = {-1, 0, 1, 0}, dy[] = {0, 1, 0, -1};
int h[N], e[M], ne[M], w[M], idx;
void add(int a, int b, int v = 0) {
	e[idx] = b, w[idx] = v, ne[idx] = h[a], h[a] = idx ++;
}
int n, m, k;
int x, y;
LL f[N][2];
int qmi(int a, int b) {
	int res = 1;
	while (b) {
		if (b & 1) res = (LL) res * a % mod;
		a = (LL) a * a % mod;
		b >>= 1;
	}
	return res;
}
void dfs(int u, int fa) {
	f[u][0] = x, f[u][1] = y;
	for (int i = h[u]; ~i; i = ne[i]) {
		int j = e[i];
		if (j == fa) continue ; 
		dfs(j, u);
		f[u][0] = f[u][0] * (f[j][0] + f[j][1]) % mod;	
		f[u][1] = f[u][1] * (f[j][0] + f[j][1] * k % mod) % mod;
	}
}
int main() {
	ios::sync_with_stdio(false), cin.tie(0);
	cin >> n >> x >> y;
	memset(h, -1, sizeof h);
	k = (LL)(y - 1) * qmi(y, mod - 2) % mod;
	for (int i = 1; i < n; i ++) {
		int u, v; cin >> u >> v;
		add(u, v), add(v, u);
	}

	dfs(1, -1);
	
	cout << (f[1][0] + f[1][1]) % mod << '\n';
	return 0;
}

因为每个点都要选一种颜色，因此$f[u,0]:以u为根且u填了一个普通颜色的方案数，f[u,1]:以u为根且u了一个填特殊颜色的方案数$，

则转移时候$f[u,0]?=f[v_i,0]?x+f[v_i,1]?y$，由于$u$选了一个因此$v_i$只能选$y?1$个特殊颜色，所以$f[u,1]?=f[v_i,0]?x+f[v_i,1]?(y?1)$

Code

#include <bits/stdc++.h>
#define x first
#define y second
#define debug(x) cout<<#x<<":"<<x<<endl;
using namespace std;
typedef long double ld;
typedef long long LL;
typedef pair<int, int> PII;
typedef pair<double, double> PDD;
typedef unsigned long long ULL;
const int N = 1e6 + 10, M = 2 * N, INF = 0x3f3f3f3f, mod = 998244353;
const double eps = 1e-8, pi = acos(-1), inf = 1e20;
int dx[] = {-1, 0, 1, 0}, dy[] = {0, 1, 0, -1};
int h[N], e[M], ne[M], w[M], idx;
void add(int a, int b, int v = 0) {
	e[idx] = b, w[idx] = v, ne[idx] = h[a], h[a] = idx ++;
}
int n, m, k;
int x, y;
LL f[N][2];
void dfs(int u, int fa) {
	f[u][0] = 1, f[u][1] = 1;
	for (int i = h[u]; ~i; i = ne[i]) {
		int j = e[i];
		if (j == fa) continue ; 
		dfs(j, u);
		f[u][0] = (f[u][0] * (f[j][0] * x % mod + f[j][1] * y % mod) % mod) % mod;	
		f[u][1] = (f[u][1] * (f[j][0] * x % mod + f[j][1] * (y - 1)  % mod) % mod) % mod;
	}
	
}
int main() {
	ios::sync_with_stdio(false), cin.tie(0);
	cin >> n >> x >> y;
	memset(h, -1, sizeof h);
	for (int i = 1; i < n; i ++) {
		int u, v; cin >> u >> v;
		add(u, v), add(v, u);
	}

	dfs(1, -1);
	
	cout << (f[1][0] * x + f[1][1] * y) % mod << '\n';
	return 0;
}