P16961 [SCCPC 2026] Eternal Augustus.

There is a tree $T$ with $n$ nodes, rooted at node $1$. Initially, node $i$ has color $c_i$ ($0 \le c_i \le 1$). Little L performed several (possibly $0$) operations. In each operation, she chooses a node $u$ such that $c_u = 0$, and for every node $v$ in the subtree of $u$, she performs $c_v \gets 1 - c_v$. After all operations, we obtain a tree $T'$, where the color of node $i$ becomes $c'_i$. Now you are given the final tree $T'$ and the color $c'_i$ of each node. You need to compute how many different possible initial states $T$ there could be. Output the answer modulo $998244353$. In this problem, we consider two trees $T_1, T_2$ to be different if and only if there exists a node $1 \le u \le n$ such that its color in $T_1$ is $c_{1,u}$, its color in $T_2$ is $c_{2,u}$, and $c_{1,u} \neq c_{2,u}$.

The first line contains a positive integer $n$ ($1 \le n \le 2 \times 10^5$), indicating the number of nodes in $T'$. The second line contains $n$ integers $c'_1, c'_2, \cdots, c'_n$ ($0 \le c'_i \le 1$), representing the color of each node in the final state. The next $n - 1$ lines each contain two positive integers $u, v$ ($1 \le u, v \le n$, $u \neq v$), indicating that there is an edge $(u, v)$ in $T'$. It is guaranteed that all edges form a tree.

Output one integer in a single line, representing the number of possible initial states $T$ modulo $998244353$.

Translated by ChatGPT 5