CF1118F2.Tree Cutting (Hard Version)

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$.

给你一棵包含 $n$ 个顶点的无向树。

其中一些顶点被染成 $k$ 种颜色之一，其余顶点未染色。保证该树中至少包含每种颜色的一个顶点。未染色的顶点可能不存在。

你选择恰好 $k - 1$ 条边构成一个子集，并将这些边从树中移除。此时树将分裂为 $k$ 个连通分量。若每个连通分量中均不包含两种或以上不同颜色的顶点，则称该边子集是“优美的”（nice）。

问：给定树中一共有多少个“优美的”边子集？若两个子集中存在某条边仅属于其中之一，则认为这两个子集不同。

答案可能很大，请对 $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.

第一行包含两个整数 $n$ 和 $k$（$2 \le n \le 3 \cdot 10^5$，$2 \le k \le n$）—— 分别表示树中顶点的数量和颜色种类数。

第二行包含 $n$ 个整数 $a_1, a_2, \dots, a_n$（$0 \le a_i \le k$）—— 表示各顶点的颜色。若 $a_i = 0$，则表示顶点 $i$ 未着色；若为其他值，则表示顶点 $i$ 着该颜色。

接下来的 $n - 1$ 行中，第 $i$ 行包含两个整数 $v_i$ 和 $u_i$（$1 \le v_i, u_i \le n$，$v_i \ne u_i$）—— 表示树的一条边。保证所给边构成一棵树。保证该树中至少包含每种颜色 $1$ 至 $k$ 的至少一个顶点。可能不存在未着色的顶点。

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$.

输出一个整数——给定树中“好”边子集的数量。若两个子集之间存在某条边在其中一个子集中出现而在另一个子集中未出现，则认为这两个子集不同。

答案可能很大，因此请对 $998244353$ 取模后输出。

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)}$.

第一个样例中的树如下所示：

唯一的“优美”子集是边 $(2, 4)$。移除该边后，树分裂为两个连通分量：$\{4\}$ 和 $\{1, 2, 3, 5\}$。其中第一个分量仅包含颜色为 $1$ 的顶点，第二个分量仅包含颜色为 $2$ 的顶点以及未着色的顶点。

第二个样例中的树如下所示：

“优美”子集有：$\{(1, 3), (4, 7)\}$、$\{(1, 3), (7, 2)\}$、$\{(3, 6), (4, 7)\}$ 和 $\{(3, 6), (7, 2)\}$。