103406: [Atcoder]ABC340 G - Leaf Color

Memory Limit:256 MB Time Limit:2 S
Judge Style:Text Compare Creator:
Submit:0 Solved:0

Description

Score: $600$ points

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

Input

Output

分数:600分

问题描述

有一个有N个顶点的树T,顶点编号从1到N。第i条边连接顶点u_i和v_i。此外,顶点i被涂成颜色A_i。

  • 找出满足以下条件的(非空)顶点子集S的数量,模998244353:
什么是诱导子图? 让S是图G的顶点子集。G由S诱导的子图是一个图,其顶点集为S,其边集由G中两端点都在S中的所有边组成。

约束条件

  • 1≤N≤2×10^5
  • 1≤A_i≤N
  • 1≤u_i<v_i≤N
  • 输入给出的图是一个树。
  • 所有输入值都是整数。

输入

输入以标准输入的以下格式给出:

$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

3
1 2 1
1 2
2 3

样例输出1

4

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

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

样例输入2

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

样例输出2

9

样例输入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

样例输出3

48

加入题单

上一题 下一题 算法标签: