题解 | #G 魔法树#

计数类树形 DP 的常用思路之一为决策边，将点或点所在连通块状态放在主状态的附加维中。例如CF461B，NOIP2022 T3都可以用这个思路解决。

回到本题，决策边的提示已经挂脸上了，考虑状态设计：注意每个点所在的连通块可以通过是否和儿子连边从而和儿子合并成一个大连通块来改变奇偶性，但是答案中整体奇偶性是必须一致的，因此二者必须都进入状态。

设 $dp_{u,i=0/1,j=0/1}$ 为，结点 $u$ 的子树中，结点 $u$ 所在连通块奇偶性为 $i$ （ $1$ 奇 $0$ 偶，下同），子树整体奇偶性为 $j$ 时子树内的方案数，初始状态显然为 $dp_{u,a_u \And 1,0/1}=1$ 。

考虑转移。由于子状态数量很小为 $4$ 个，考虑枚举 $16$ 种情况进行转移，这是最为简单粗暴且不容易出错的办法：

以下设 $v\in son_u$ ，即 $v$ 为 $u$ 的儿子：

$dp_{u,0,0},dp_{v,0,0}$ ：二者连边时，转移到 $dp_{u,0,0}$ ，不连边时仍然转移到 $dp_{u,0,0}$ ；

$dp_{u,0,0},dp_{v,0,1}$ ：无论是否连边，都会出现整个子树除了 $u$ 所在连通块外仍有异色连通块的情况，不合法；同理， $dp_{u,0,0},dp_{v,1,1}$ ， $dp_{u,0,1},dp_{v,0,0}$ ， $dp_{u,0,1},dp_{v,1,0}$ ， $dp_{u,1,0},dp_{v,0,1}$ ， $dp_{u,1,0},dp_{v,1,1}$ ， $dp_{u,1,1},dp_{v,0,0}$ ， $dp_{u,1,1},dp_{v,1,0}$ 的组合无论是否连边都不合法，之后将忽略对她们的讨论；

$dp_{u,0,0},dp_{v,1,0}$ ：二者连边时，转移到 $dp_{u,1,0}$ ；不连边时，出现除子树根所在连通块之外的异色连通块，不合法。

$dp_{u,0,1},dp_{v,0,1}$ ：二者连边时，转移到 $dp_{u,0,1}$ ；不连边时，出现除子树根所在连通块之外的异色连通块，不合法。

$dp_{u,0,1},dp_{v,1,1}$ ：二者连边时，转移到 $dp_{u,1,1}$ ；不连边时，转移到 $dp_{u,0,1}$ 。

$dp_{u,1,0},dp_{v,0,0}$ ：二者连边时，转移到 $dp_{u,1,0}$ ；不连边时，仍然转移到 $dp_{u,1,0}$ 。

$dp_{u,1,0},dp_{v,1,0}$ ：二者连边时，转移到 $dp_{u,0,0}$ ；不连边时，出现除子树根所在连通块之外的异色连通块，不合法。

$dp_{u,1,1},dp_{v,0,1}$ ：二者连边时，转移到 $dp_{u,1,1}$ ；不连边时，出现除子树根所在连通块之外的异色连通块，不合法。

$dp_{u,1,1},dp_{v,1,1}$ ：二者连边时，转移到 $dp_{u,0,1}$ ；不连边时，转移到 $dp_{u,1,1}$ 。

综上，转移方程为：

$dp_{u,0,0}\leftarrow dp_{u,1,0}\times dp_{v,1,0} 2\times dp_{u,0,0}\times dp_{u,0,0}$

$dp_{u,0,1}\leftarrow dp_{u,1,1}\times dp_{v,1,1} dp_{u,0,1}\times dp_{v,1,1} dp_{u,0,1}\times dp_{v,0,1}$

$dp_{u,1,0}\leftarrow 2\times dp_{u,1,0}\times dp_{v,0,0} dp_{u,0,0}\times dp_{v,1,0}$

$dp_{u,1,1}\leftarrow dp_{u,1,1}\times dp_{v,1,1} dp_{u,1,1}\times dp_{v,0,1} dp_{u,0,1}\times dp_{v,1,1}$

注意转移开始之前先复制原来前继状态的值防止使用更新之后的值转移导致出错。

根据状态定义，答案为 $dp_{root,0,0} dp_{root,1,1}$ 。

注意取模，时间复杂度线性。

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+11;
using ll=long long;
const ll p=998244353;
struct edge
{
	int t,nxt;
}G[N*2];int at[N],ct;
inline void add_edge(int u,int v)
{
	G[++ct]={v,at[u]};at[u]=ct;
}
ll dp[N][2][2];int a[N];
//dp[u][0][0]=dp[u][0][0]*2*dp[v][0][0]+dp[u][1][0]*dp[v][1][0]
//dp[u][0][1]=dp[u][0][1]*dp[v][0][1]+dp[u][0][1]*dp[v][1][1]+dp[u][1][1]*dp[v][1][1]
//dp[u][1][0]=dp[u][1][0]*2*dp[v][0][0]+dp[u][0][0]*dp[v][1][0]
//dp[u][1][1]=dp[u][0][1]*dp[v][1][1]+dp[u][1][1]*dp[v][0][1]+dp[u][1][1]*dp[v][1][1]
void treedp(int u,int f)
{
	dp[u][a[u]&1][0]=dp[u][a[u]&1][1]=1;
	//if(at[u]) dp[u][a[u]&1][!a[u]&1]=1;
	for(int i=at[u];i;i=G[i].nxt)
	{
		auto v=G[i].t;if(v==f) continue;
		treedp(v,u);
		ll x=dp[u][0][0],y=dp[u][0][1],z=dp[u][1][0],w=dp[u][1][1];
		dp[u][0][0]=((2ll*x*dp[v][0][0]%p)+(z*dp[v][1][0]%p))%p;
		dp[u][0][1]=((y*dp[v][0][1]%p)+(y*dp[v][1][1]%p)+(w*dp[v][1][1]%p))%p;
		dp[u][1][0]=((2ll*z*dp[v][0][0]%p)+(x*dp[v][1][0]%p))%p;
		dp[u][1][1]=((y*dp[v][1][1]%p)+(w*dp[v][0][1]%p)+(w*dp[v][1][1]%p))%p;
	}
}
int n,u,v;
int main()
{
	cin>>n;
	for(int i=1;i<=n;++i) cin>>a[i];
	for(int i=1;i<n;++i) cin>>u>>v,add_edge(u,v),add_edge(v,u);
	treedp(1,0);
	cout<<(dp[1][0][0]+dp[1][1][1])%p;//<<" "<<dp[1][0][0]<<" "<<dp[1][0][1]<<" "<<dp[1][1][0]<<" "<<dp[1][1][1];
	return 0;
}