题解：P7727 风暴之眼（Eye of the Storm）

思路

由题意可以很直观地想到，直接根据 $4$ 种初始状态设计转移，分别对应 $(1,\texttt{OR})$ 、 $(0,\texttt{OR})$ 、 $(1,\texttt{AND})$ 、 $(0,\texttt{AND})$ 。

但是这样设计时，对于目标值为 $0$ 和 $1$ 的结点需要分别设计转移，比较复杂。

把最终目标值一并考虑后，可以将这些情况统一归纳为 $5$ 个状态（其中包含 $1$ 个非法状态）。这里的状态指的是初始值相对于目标值所扮演的角色。

脆弱态

结点当前已经是目标值，但自身不能保证一直维持该值。只要受到相反权值的邻居影响（后文称作污染），就会偏离目标值。
- 当目标值 $a_u=0$ 时：
  
  $(w_u,t_u)=(0,\texttt{OR})$
  
  此时只要某个邻居出现 $1$ ，结点 $u$ 就会变为 $1$ 。
- 当目标值 $a_u=1$ 时：
  
  $(w_u,t_u)=(1,\texttt{AND})$
  
  此时只要某个邻居出现 $0$ ，结点 $u$ 就会变为 $0$ 。
因此，状态 0 的结点不能处于异色连通块的边界，其所有邻居的目标值都必须与它相同。即使邻居的目标值相同，也必须保证其初始值不会在传播过程中污染 $u$ ，这一限制会体现在后面的转移中。
渴求态

结点的初始值与目标值不同，需要某个邻居向它传递正确的权值，但其子树中暂时不存在这样的邻居，只能依靠父亲方向将它转变为目标值。
- 当目标值 $a_u=0$ 时：
  
  $(w_u,t_u)=(1,\texttt{AND})$
  
  结点 $u$ 需要某个邻居变为 $0$ ，才能通过 $\texttt{AND}$ 运算变为 $0$ 。
- 当目标值 $a_u=1$ 时：
  
  $(w_u,t_u)=(0,\texttt{OR})$
  
  结点 $u$ 需要某个邻居变为 $1$ ，才能通过 $\texttt{OR}$ 运算变为 $1$ 。
满足态

状态 2 与状态 1 对应的初始值和类型相同，区别在于其子树中已经存在一个能够向它传递目标值的结点，因此不再需要父亲提供帮助。

一旦结点变为目标值，由于它的运算包含自身，它之后便不会再次偏离目标值。
- 当目标值 $a_u=0$ 时：
  
  $(w_u,t_u)=(1,\texttt{AND})$
  
  子树中已经存在能够向 $u$ 传递 $0$ 的结点。
- 当目标值 $a_u=1$ 时：
  
  $(w_u,t_u)=(0,\texttt{OR})$
  
  子树中已经存在能够向 $u$ 传递 $1$ 的结点。
强源态

结点初始时就是目标值，并且无论邻居如何变化，它都不会偏离目标值，还可以持续向相邻结点传播正确的权值。
- 当目标值 $a_u=0$ 时：
  
  $(w_u,t_u)=(0,\texttt{AND})$
  
  由于 $\texttt{AND}$ 运算包含自身，而自身已经为 $0$ ，所以它之后始终为 $0$ 。
- 当目标值 $a_u=1$ 时：
  
  $(w_u,t_u)=(1,\texttt{OR})$
  
  由于 $\texttt{OR}$ 运算包含自身，而自身已经为 $1$ ，所以它之后始终为 $1$ 。
非法态（不计入 DP）

对于不同的目标值，还各存在一种不可能达到目标值的初始状态，因此不计入任何 DP 状态。
- 当目标值 $a_u=0$ 时， $(w_u,t_u)=(1,\texttt{OR})$ 非法。
  
  由于运算包含结点自身，而 $u$ 的初始值为 $1$ ，所以每次进行 $\texttt{OR}$ 运算后它都仍然为 $1$ ，永远不可能变为目标值 $0$ 。
- 当目标值 $a_u=1$ 时， $(w_u,t_u)=(0,\texttt{AND})$ 非法。
  
  由于运算包含结点自身，而 $u$ 的初始值为 $0$ ，所以每次进行 $\texttt{AND}$ 运算后它都仍然为 $0$ ，永远不可能变为目标值 $1$ 。

故可以设 $dp[u][0\ldots 3]$ 表示以 $u$ 为根的子树中，使得 $u$ 处于对应状态的合法方案数。

在转移过程中，我们按顺序将结点 $u$ 的各个儿子合并进来。设当前正在合并的儿子为 $v$ ：

$u_i$ ：在合并儿子 $v$ 之前，结点 $u$ 已经处理完的部分中， $u$ 处于状态 $i$ 的方案数。
$v_i$ ：以结点 $v$ 为根的完整子树中， $v$ 处于状态 $i$ 的方案数。
$n_i$ ：将儿子 $v$ 合并进来之后，结点 $u$ 处于状态 $i$ 的新方案数。

这里的下标 $i\in\{0,1,2,3\}$ 对应前文定义的四种状态。

每次完成转移后，将新值赋回：

$u_i\leftarrow n_i\qquad(0\le i<4)$

转移需要根据边 $(u,v)$ 两端的最终目标 $a_u$ 与 $a_v$ 是否相同分类讨论。

同色边

当 $a_u=a_v$ 时，两端的目标值相同，正确权值可以沿着同色连通块逐步传播。

$n_0=u_0\times(v_0+v_3)$ ：

$u$ 处于脆弱态，因此 $v$ 必须从一开始就保持与目标值相同。合法的状态只有同样脆弱的 $v_0$ 和稳定的强源态 $v_3$ 。
$n_1=u_1\times v_1$ ：

$u$ 尚未被满足，因此 $v$ 也不能向 $u$ 提供目标值，只能处于状态 $v_1$ 。

当 $u$ 之后从父亲方向获得满足时，正确权值还可以继续由 $u$ 传递给 $v$ 。
$n_2=u_1\times(v_2+v_3)+u_2\times(v_1+v_2+v_3)$ ：

$u$ 被满足有两种途径。

第一种是 $u$ 原本处于渴求态 $u_1$ ，并从已经满足的 $v_2$ 或强源态 $v_3$ 处获得目标值。

第二种是 $u$ 原本已经处于满足态 $u_2$ 。此时 $v$ 只要不是脆弱态 $v_0$ 即可，因为 $u_2$ 的初始值与目标值相反，会在第一秒污染 $v_0$ 。
$n_3=u_3\times(v_0+v_1+v_2+v_3)$ ：

$u$ 是稳定的强源态，不会被 $v$ 改变，并且还能向需要满足的 $v_1$ 传播正确的权值，因此 $v$ 可以处于任意状态。

异色边

当 $a_u\ne a_v$ 时，两端的最终目标值相反。

异色边既要考虑两端初始值在第一秒产生的影响，也要保证结点稳定后不会继续污染对方的脆弱态。

由于 $a_u\ne a_v$ ，所以 $v_1$ 和 $v_2$ 的初始值与 $a_v$ 相反，恰好等于 $a_u$ ；而 $v_0$ 和 $v_3$ 的初始值等于 $a_v$ ，与 $a_u$ 相反。

$n_0=0$ ：

$v$ 最终会稳定为 $a_v$ ，而 $a_v\ne a_u$ 。因此 $u$ 迟早会受到相反权值的污染，脆弱态一定非法。
$n_1=u_1\times v_3$ ：

若 $v$ 处于 $v_1$ 或 $v_2$ ，其初始值恰好等于 $a_u$ ，会直接满足 $u$ ，结果应转入状态 2。

若 $v$ 处于脆弱态 $v_0$ ，当 $u$ 之后从父亲处获得满足并变为 $a_u$ 时，会反过来污染 $v_0$ 。

因此，只有稳定且不会满足 $u$ 的强源态 $v_3$ 合法。
$n_2=u_1\times(v_1+v_2)+u_2\times(v_1+v_2+v_3)$ ：

当 $u$ 原本处于渴求态 $u_1$ 时， $v_1$ 和 $v_2$ 的初始值恰好等于 $a_u$ ，可以在第一秒满足 $u$ 。

当 $u$ 原本已经处于满足态 $u_2$ 时， $v_0$ 会在 $u$ 变为目标值后受到污染，因此不合法； $v_1$ 、 $v_2$ 和 $v_3$ 均可以与 $u$ 共存。
$n_3=u_3\times(v_2+v_3)$ ：

$u_3$ 的初始值和最终值均为 $a_u$ ，与 $a_v$ 相反。

若 $v$ 处于脆弱态 $v_0$ ，会立即被 $u_3$ 污染；若 $v$ 处于渴求态 $v_1$ ，其子树无法满足它，而 $u_3$ 也无法向它提供目标值 $a_v$ 。

因此， $v$ 只能处于已经被子树满足的 $v_2$ 或稳定的强源态 $v_3$ 。

所有儿子合并完成后，若根结点仍处于状态 1，它已经没有父亲可以依赖，因此这种方案不合法。最终答案为：

$dp[1][0]+dp[1][2]+dp[1][3]$

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N = 2e5 + 5;
const int mod = 998244353;

int n, a[N];
vector<int> g[N];
bool b[N];
ll dp[N][4];

void dfs(int u, int p)
{
    // 异色边界上的结点不能处于状态 0
    dp[u][0] = b[u] ? 0 : 1;
    dp[u][1] = 1;
    dp[u][2] = 0;//表示已经被子树满足，一开始还没有子树信息，因此初始方案数为 $0$
    dp[u][3] = 1;

    for (int v : g[u])
    {
        if (v == p)
            continue;

        dfs(v, u);

        ll v0 = dp[v][0], v1 = dp[v][1];
        ll v2 = dp[v][2], v3 = dp[v][3];
        ll n0, n1, n2, n3;

        if (a[u] == a[v])
        {
            n0 = dp[u][0] * (v0 + v3) % mod;
            n1 = dp[u][1] * v1 % mod;
            n2 = (dp[u][1] * (v2 + v3) % mod
                + dp[u][2] * (v1 + v2 + v3) % mod) % mod;
            n3 = dp[u][3] * (v0 + v1 + v2 + v3) % mod;
        }
        else
        {
            n0 = 0;
            n1 = dp[u][1] * v3 % mod;
            n2 = (dp[u][1] * (v1 + v2) % mod
                + dp[u][2] * (v1 + v2 + v3) % mod) % mod;
            n3 = dp[u][3] * (v2 + v3) % mod;
        }

        dp[u][0] = n0;
        dp[u][1] = n1;
        dp[u][2] = n2;
        dp[u][3] = n3;
    }
}

int main()
{
    cin.tie(0)->sync_with_stdio(0);

    cin >> n;
    for (int i = 1; i <= n; i++)
        cin >> a[i];

    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    // 标记异色连通块的边界
    for (int u = 1; u <= n; u++)
    {
        for (int v : g[u])
        {
            if (a[u] != a[v])
            {
                b[u] = true;
                break;
            }
        }
    }

    dfs(1, 0);

    // 根结点不能处于尚未满足的状态 1
    ll ans = (dp[1][0] + dp[1][2] + dp[1][3]) % mod;
    cout << ans << "\n";

    return 0;
}