树上DP

题目链接：
https://ac.nowcoder.com/acm/contest/11187/D

题意：
给出一棵 n 个点的树，有 x 种普通颜色，y 种特殊颜色现在要给树上的每个节点染色，
普通颜色染色没有限制，但两个相邻的节点不能染相同颜色的特殊颜色，求染色方案数，答案对 998244353 取模。

解答：
树上DP。从根节点遍历到叶子节点，dp[u][1/0] 来保存以u为根，填上普通颜色（或者特殊颜色）的方案数。

#include <iostream>
#include <vector>
using namespace std;
#define int long long
const int N = 1e6+5;
const int mod = 998244353;
int dp[N][2], x, y;
vector<int> vis[N];
 
inline int read() { // 内联函数  调用时直接将整个代码复制过去，减少了调用函数的时间 
	int f = 1, k = 0;
	char c = getchar();
	while (c<'0' || c>'9') {if (c=='-') f = -1; c = getchar();}
	while (c>='0' && c<='9') {k = k*10+c-'0'; c = getchar();}
	return f*k;
}

void dfs(int u, int fa) { // 由根节点 逐步 往下遍历 
    dp[u][0] = 1, dp[u][1] = 1; // 以u为根 且u填了一个普通颜色（特殊颜色）时，子树的染色方案数   均初始化为 1 
	for (int i=0; i<vis[u].size(); i++) {
		if (vis[u][i] == fa) continue; // 若 子节点 等于 其父节点 则 continue 
		dfs(vis[u][i], u); // 继续往下遍历
		dp[u][0] = dp[u][0]*(dp[vis[u][i]][0]*x%mod+dp[vis[u][i]][1]*y%mod)%mod; // 该节点 * 子节点可能的颜色
		dp[u][1] = dp[u][1]*(dp[vis[u][i]][0]*x%mod+dp[vis[u][i]][1]*(y-1)%mod)%mod; // 因为特殊颜色不能相邻 故其子节点特殊颜色种数 y-1  
	}
	return ;
}

signed main() { // signed = signed int = int
	int n, u, v;
	while (cin >> n >> x >> y) {
		for (int i=1; i<=n-1; i++) {
			u = read(), v = read();
			vis[u].push_back(v); // 利用 vector数组 来建树 
			vis[v].push_back(u);
		}
		dfs(1, 0);
		cout << (dp[1][0]*x%mod+dp[1][1]*y%mod)%mod << endl; // 根节点 分别乘上 x 种普通颜色和 y 种特殊颜色，即得到总的方案数
	}
	return 0;
}