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 翻译