AT_abc340_g [ABC340G] Leaf Color

有一棵包含 $N$ 个顶点的树 $T$，顶点编号为 $1$ 到 $N$。第 $i$ 条边连接顶点 $u_i$ 和顶点 $v_i$。此外，顶点 $i$ 被染成颜色 $A_i$。 请计算满足以下条件的 $T$ 的顶点集合（非空子集）$S$ 的个数，并对 $998244353$ 取模： - $S$ 所对应的诱导子图 $G$ 满足以下所有条件： - $G$ 是一棵树。 - 所有度数为 $1$ 的顶点的颜色都相同。 诱导子图的定义如下：对于图 $G$ 的顶点子集 $S$，$S$ 所对应的诱导子图是指顶点集合为 $S$，边集合为“$G$ 中两端都属于 $S$ 的所有边”的图。

输入以如下格式从标准输入读入： > $N$ > $A_1\ A_2\ \dots\ A_N$ > $u_1\ v_1$ > $u_2\ v_2$ > $\vdots$ > $u_{N-1}\ v_{N-1}$

输出满足题目条件的顶点集合 $S$ 的个数，对 $998244353$ 取模。

## 限制条件 - $1 \leq N \leq 2 \times 10^5$ - $1 \leq A_i \leq N$ - $1 \leq u_i < v_i \leq N$ - 输入保证图为一棵树 - 输入的所有数均为整数 ## 样例解释 1 满足条件的顶点集合有以下 $4$ 种： - $\lbrace 1 \rbrace$ - $\lbrace 1, 2, 3 \rbrace$ - $\lbrace 2 \rbrace$ - $\lbrace 3 \rbrace$ 由 ChatGPT 4.1 翻译