Tree Cutting (Hard Version)

题意翻译

        给定一个有$n$个节点的树, 结点可能有颜色, 共$k$种颜色, 颜色编号$1...k$，每种颜色都出现。有的点没有颜色, 用$0$表示. 将其删去$k-1$条边, 即划分成$k$个联通块, 使每个联通块中**恰好**含一种颜色, 颜色为$0$的节点可以在任意联通块中. 求划分的方案数. ( 无解输出$0$, 答案对$998244353$取模. )

题目描述

You are given an undirected tree of $ n $ vertices. Some vertices are colored one of the $ k $ colors, some are uncolored. It is guaranteed that the tree contains at least one vertex of each of the $ k $ colors. There might be no uncolored vertices. You choose a subset of exactly $ k - 1 $ edges and remove it from the tree. Tree falls apart into $ k $ connected components. Let's call this subset of edges nice if none of the resulting components contain vertices of different colors. How many nice subsets of edges are there in the given tree? Two subsets are considered different if there is some edge that is present in one subset and absent in the other. The answer may be large, so print it modulo $ 998244353 $ .

输入输出格式

输入格式

The first line contains two integers $ n $ and $ k $ ( $ 2 \le n \le 3 \cdot 10^5 $ , $ 2 \le k \le n $ ) — the number of vertices in the tree and the number of colors, respectively. The second line contains $ n $ integers $ a_1, a_2, \dots, a_n $ ( $ 0 \le a_i \le k $ ) — the colors of the vertices. $ a_i = 0 $ means that vertex $ i $ is uncolored, any other value means the vertex $ i $ is colored that color. The $ i $ -th of the next $ n - 1 $ lines contains two integers $ v_i $ and $ u_i $ ( $ 1 \le v_i, u_i \le n $ , $ v_i \ne u_i $ ) — the edges of the tree. It is guaranteed that the given edges form a tree. It is guaranteed that the tree contains at least one vertex of each of the $ k $ colors. There might be no uncolored vertices.

输出格式

Print a single integer — the number of nice subsets of edges in the given tree. Two subsets are considered different if there is some edge that is present in one subset and absent in the other. The answer may be large, so print it modulo $ 998244353 $ .

输入输出样例

输入样例 #1

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

输出样例 #1

1

输入样例 #2

7 3
0 1 0 2 2 3 0
1 3
1 4
1 5
2 7
3 6
4 7

输出样例 #2

4

说明

Here is the tree from the first example: The only nice subset is edge $ (2, 4) $ . Removing it makes the tree fall apart into components $ \{4\} $ and $ \{1, 2, 3, 5\} $ . The first component only includes a vertex of color $ 1 $ and the second component includes only vertices of color $ 2 $ and uncolored vertices. Here is the tree from the second example: The nice subsets are $ \{(1, 3), (4, 7)\} $ , $ \{(1, 3), (7, 2)\} $ , $ \{(3, 6), (4, 7)\} $ and $ \{(3, 6), (7, 2)\} $ .