Score: $600$ points

问题陈述

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

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

• 由 $S$ 确定的 $T$ 的子图 $G$ 满足所有以下条件：
  • $G$ 是一棵树。
  • 所有度数为 $1$ 的顶点颜色相同。

什么是导出子图？设 $S$ 是图 $G$ 中顶点的一个子集。由 $S$ 导出的子图是一个图，其顶点集合为 $S$，边集合包含 $G$ 中两端点均在 $S$ 内的所有边。

以上为通义千问 qwen-max 翻译，仅供参考。

Problem Statement

There is a tree $T$ with $N$ vertices numbered from $1$ to $N$. The $i$-th edge connects vertices $u_i$ and $v_i$. Additionally, vertex $i$ is painted with color $A_i$.
Find the number, modulo $998244353$, of (non-empty) subsets $S$ of the vertex set of $T$ that satisfy the following condition:

• The induced subgraph $G$ of $T$ by $S$ satisfies all of the following conditions:
  • $G$ is a tree.
  • All vertices with degree $1$ have the same color.

What is an induced subgraph? Let $S$ be a subset of the vertices of a graph $G$. The induced subgraph of $G$ by $S$ is a graph whose vertex set is $S$ and whose edge set consists of all edges in $G$ that have both endpoints in $S$.

Constraints

• $1 \leq N \leq 2 \times 10^5$
• $1 \leq A_i \leq N$
• $1 \leq u_i \lt v_i \leq N$
• The graph given in the input is a tree.
• All input values are integers.

Input

The input is given from Standard Input in the following format:

$N$

$A_1$ $A_2$ $\dots$ $A_N$

$u_1$ $v_1$

$u_2$ $v_2$

$\vdots$

$u_{N-1}$ $v_{N-1}$

Output

Print the number, modulo $998244353$, of (non-empty) subsets $S$ of the vertex set of $T$ that satisfy the condition in the problem statement.

Sample Input 1

3
1 2 1
1 2
2 3

Sample Output 1

4

The following four sets of vertices satisfy the condition.

• $\lbrace 1 \rbrace$
• $\lbrace 1, 2, 3 \rbrace$
• $\lbrace 2 \rbrace$
• $\lbrace 3 \rbrace$

Sample Input 2

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

Sample Output 2

9

Sample Input 3

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

Sample Output 3

48

update @ 2024/3/10 01:33:43