G - Leaf Color Editorial /

Time Limit: 2 sec / Memory Limit: 1024 MB

配点 : 600 点

問題文

頂点に 1 から N までの番号がついた N 頂点の木 T があります。i 番目の辺は頂点 u_i と頂点 v_i を結んでいます。また、頂点 i は色 A_i で塗られています。
T の頂点集合の(非空な)部分集合 S のうち、次の条件を満たすものの個数を 998244353 で割った余りを求めてください。

  • T の S による誘導部分グラフ G は次の条件を全て満たす。
    • G は木である。
    • 次数 1 の頂点に塗られた色が全て一致している。
誘導部分グラフとは S をグラフ G の頂点の部分集合とします。このとき、G の S による誘導部分グラフとは、頂点集合が S で、辺集合が「G の辺であって両端が S に含まれるもの全て」であるようなグラフです。

制約

  • 1 \leq N \leq 2 \times 10^5
  • 1 \leq A_i \leq N
  • 1 \leq u_i \lt v_i \leq 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

条件を満たす頂点の集合は次の 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

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