牛客11187D - 月之暗面

链接：https://ac.nowcoder.com/acm/contest/11187/D 知识点：树形DP、组合计数 难度：绿

题意

• 给出一棵 n 个点的树， 有 X 种普通颜色，Y 种特殊颜色
• 现在要给树上的每个节点染色，普通颜色染色没有限制， 但两个相邻的节点不能染相同颜色的特殊颜色 这些颜色不要求全部使用 求染色方案数， 答案对 998244353 取模。

思路

考虑树形DP

错误思路

• 用 $dp[cur][0\mathrm{/}1]dp[cur][0/1]$ 表示 $curcur$ 节点染普通颜色和特殊颜色的方案数。

• 假设该父节点要染特殊颜色，并且父节点的仅某一个儿子节点染了特殊颜色， 那么这个父节点能染 $Y-1Y-1$ 种颜色。 但是如果父节点的染了特殊颜色的儿子有多个呢？很难处理。

• 但是，对于一个节点，它的父节点却只有一个。这就有了下面的正确思路。

正确思路

• 一般的树形DP都是自下向上做计算，这次考虑自上向下做计算。

• 我们先染父节点。 假设父节点要染特殊颜色，父节点有 Y 种方案， 假设在父节点染特殊颜色的基础上，它的儿子节点也要染特殊颜色，则儿子节点就有 Y-1 种方案。 假设在父节点染普通颜色的基础上，它的儿子节点要染特殊颜色，则儿子节点就有 Y 种方案。

• 我们发现，染色既和当前节点有关，又和父节点有关。

• 用 $dp[cur][0\mathrm{/}1][0\mathrm{/}1]dp[cur][0/1][0/1]$ 表示 $curcur$ 节点染普通颜色还是特殊颜色，并且它的父亲节点染普通颜色还是特殊颜色的方案数。

代码

#include <cstdio>
#include <iostream>
#include <vector>
#define int long long
const int N		= 1e6+10;
const int MOD	= 998244353;
using namespace std;

vector<int> G[N];
int dp[N][2][2];
int n, X, Y;

void DFS(int cur,int pre)
{
	dp[cur][0][0] = X;
	dp[cur][0][1] = X;
	dp[cur][1][0] = Y;
	dp[cur][1][1] = Y-1;
	
	for (auto nxt : G[cur]) 
	{
		if(nxt==pre)continue;
		DFS(nxt, cur);
		
		dp[cur][0][0] *= (dp[nxt][0][0] + dp[nxt][1][0])%MOD, dp[cur][0][0]%=MOD;
		dp[cur][0][1] *= (dp[nxt][0][0] + dp[nxt][1][0])%MOD, dp[cur][0][1]%=MOD;
		
		dp[cur][1][0] *= (dp[nxt][0][1] + dp[nxt][1][1])%MOD,dp[cur][1][0]%=MOD;
		dp[cur][1][1] *= (dp[nxt][0][1] + dp[nxt][1][1])%MOD,dp[cur][1][1]%=MOD;
	}
}

void Solve()
{
	DFS(1, 0);
	printf("%lld\n",(dp[1][0][0]+dp[1][1][0])%MOD);
}

signed main()
{
	scanf("%lld %lld %lld",&n,&X,&Y);
	for (int i=1; i<=n-1; i++)
	{
		int u,v;
		scanf("%lld %lld",&u,&v);
		G[u].push_back(v);
		G[v].push_back(u);
	}
	Solve();
	
	return 0;
}