P16961 [SCCPC 2026] 永恒的奥古斯都

有一棵 $n$ 个点的树 $T$，树的根节点为 $1$。初始时点 $i$ 有颜色 $c_i$（$0 \le c_i \le 1$）。 小 L 进行了若干次（可以为 $0$）操作，每次操作她会选择一个满足 $c_u =0$ 的节点 $u$，对 $u$ 子树内的所有点 $v$ 执行 $c_v \gets 1 - c_v$。所有操作结束后得到树 $T'$，其中点 $i$ 的颜色变为了 $c'_i$。 现在给你最后得到的树 $T'$ 和每个点的颜色 $c'_i$，你需要求出有多少种不同的可能初始状态 $T$。答案对 $998244353$ 取模。 在此题中，我们认为两棵树 $T_1,T_2$ 不同，当且仅当存在点 $1 \le u \le n$，满足其在 $T_1$ 中的颜色为 $c_{1,u}$，在 $T_2$ 中的颜色为 $c_{2,u}$，且 $c_{1,u} \neq c_{2,u}$。

输入第一行一个正整数 $n$（$1 \le n \le 2 \times 10^5$），表示 $T'$ 的点数。 第二行 $n$ 个整数 $c'_1,c'_2,\cdots,c'_n$（$0 \le c'_i \le 1$），表示最终状态下每个点的颜色。 接下来 $n-1$ 行，每行两个正整数 $u,v$（$1 \le u,v \le n$，$u \neq v$），表示 $T'$ 中存在一条边 $(u,v)$。保证所有边构成一棵树。

输出一行一个整数，表示可能的初始状态 $T$ 的个数对 $998244353$ 取模后的值。