CF1118F2 Tree Cutting (Hard Version)

给定一棵包含 $n$ 个顶点的无向树。 有些顶点被染成 $k$ 种颜色中的一种，有些顶点未染色。保证树中每种颜色至少有一个顶点。也可能没有未染色的顶点。 你需要选择恰好 $k-1$ 条边并将其从树中删除。此时树会分裂成 $k$ 个连通分量。我们称这样的边集为“优美子集”，如果删除这些边后，任意一个连通分量内的所有顶点要么未染色，要么都染成同一种颜色（即没有分量内包含不同颜色的顶点）。 问在给定的树中，有多少个“优美子集”？如果两个子集存在某条边只在其中一个子集中出现，则认为它们不同。 答案可能很大，请输出答案对 $998244353$ 取模后的结果。

第一行包含两个整数 $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$ 行，每行包含两个整数 $v_i$ 和 $u_i$（$1 \le v_i, u_i \le n$，$v_i \ne u_i$），表示树中的一条边。保证这些边构成一棵树。保证树中每种颜色至少有一个顶点。也可能没有未染色的顶点。

输出一个整数，表示“优美子集”的数量。如果两个子集存在某条边只在其中一个子集中出现，则认为它们不同。 答案对 $998244353$ 取模。

以下是第一个样例的树结构：  唯一的优美子集是边 $(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)\}$。 由 ChatGPT 4.1 翻译