[nk] 牛客练习赛97 D.月之暗面 树形dp记录方案数

前言

传送门 :

题意

思路

一眼题了属于是,第一次碰到树形$dp$ 记录方案数记录一下

状态表示 :
$dp[i][0]$表示当前节点使用普通颜色
$dp[i][1]$表示当前节点使用特殊颜色

状态转移 :
同上司的舞会,
$dp[u][0]=dp[u][0]\ast (dp[x][0]+dp[x][1]\ast y\%mod)\%mod;$
$dp[u][1]=dp[u][1]\ast (dp[x][0]+dp[x][1]\ast (y-1)\%mod)\%mod;$

Mycode

#include <iostream>
#include <vector>
#include <map>
#include <cstring>
#include <queue>
#include <set>
#include <algorithm>
using namespace std;
#define IOS  ios::sync_with_stdio(false);
#define CIT  cin.tie(0);
#define COT  cout.tie(0);

#define ll long long
#define x first
#define y second
#define pb push_back
#define endl '\n'
#define all(x) (x).begin(),x.end()
#define Fup(i,a,b) for(int i=a;i<=b;i++)
#define Fde(i,a,b) for(int i=a;i>=b;i--)

typedef priority_queue<int,vector<int>,greater<int>>  Pri_m;
typedef pair<int,int> pii;
typedef vector<int> VI;
map<int,int> mp;
const int N = 1e6+10 ,mod = 998244353;
ll  dp[N][2];
int has_fa[N];

int n,x,y;
vector<int> g[N];

void dfs(int u,int fa){
	dp[u][0] = x;
	dp[u][1] = 1;

	for(auto x  : g[u]){
		if(x == fa) continue;
		dfs(x,fa);
		dp[u][0] =  dp[u][0]*(dp[x][0]+dp[x][1]*y%mod)%mod;
		dp[u][1] =  dp[u][1]*(dp[x][0]+dp[x][1]*(y-1)%mod)%mod;
	}
}
void solve(){
	cin>>n>>x>>y;
	for(int i = 1; i<n;i++){
		int a,b;cin>>a>>b;
		g[a].pb(b);
		has_fa[b] = 1;
	}

	int root = 1;
	while(has_fa[root]) root++;

	dfs(root,0);
	cout<<(dp[root][1]*y + dp[root][0])%mod<<endl;

}

int main(){
    //int t;cin>>t;while(t--)
    solve();
    return 0 ;
}