月之暗面---树形DP+数学

题目链接

$f[u][0]$表示u节点染普通颜色
$f[u][1]$表示u节点染特殊颜色

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int mod = 998244353;
int qmi(int a,int b){int res=1;while(b){if(b&1)res=res*a%mod;a=a*a%mod;b>>=1;}return res;}
signed main()
{
	int n,x,y;cin>>n>>x>>y;
    vector<vector<int>>f(n+1,vector<int>(2)),g(n+1);
    for(int i=1,a,b;i<n;i++)
    {
    	cin>>a>>b;
    	g[a].push_back(b);
    	g[b].push_back(a);
	}
	function<void(int,int)> dfs = [&](int u,int fa)
	{
		f[u][0]=x;
		f[u][1]=y;
		for(auto v : g[u])
		{
			if(v==fa)continue;
			dfs(v,u);
			f[u][0]=f[u][0]*((f[v][0]+f[v][1])%mod)%mod; 
			f[u][1]=f[u][1]*((f[v][0]+f[v][1]*qmi(y,mod-2)%mod*(y-1))%mod)%mod;
		}
	};
	dfs(1,-1);
	int res=(f[1][0]+f[1][1])%mod;
	cout<<res<<'\n';
}