叶子的颜色

题目描述

有一棵具有 $N$ 个顶点的树 $T$，顶点编号为 $1$ 到 $N$。第 $i$ 条边连接顶点 $u_i$ 和 $v_i$。此外，顶点 $i$ 被涂上了颜色 $A_i$。

求满足以下条件的（非空）顶点子集 $S$ 的数量，对 $998244353$ 取模：

由 $S$ 生成的 $T$ 的诱导子图 $G$ 满足以下所有条件：

• $G$ 是一棵树。
• 所有度为 $1$ 的顶点颜色相同。

什么是诱导子图？

设 $S$ 是图 $G$ 的顶点集合的子集。图 $G$ 关于 $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$ 取模的结果。

样例

3
1 2 1
1 2
2 3

4

以下四个顶点集合满足条件。

$\{1\}$

$\{1, 2, 3\}$

$\{2\}$

$\{3\}$

5
2 2 1 1 1
2 5
3 4
1 3
1 5

9

15
5 3 5 1 1 4 4 4 2 5 5 4 4 2 5
3 13
4 10
7 11
8 9
2 10
2 14
5 11
5 6
6 13
12 13
9 14
9 13
1 13
1 15

48

数据范围

• $1 \le N \le 2 \times 10^5$
• $1 \le A_i \le N$
• $1 \le u_i \lt v_i \le N$
• 输入给出的图是一棵树
• 所有输入值均为整数