CodeForces 1118F2. Tree Cutting (Hard Version)

题目简述：给定$n \leq 3 \times 10^5$个节点的树，其中一部分节点被染色，一共有$k$种不同的颜色。求将树划分成 $k$ 个不相交的部分的方案数，使得每个部分中除了未染色的节点以外的所有节点颜色相同，答案模$998244353$（质数）。

解：code

Step 1. 缩点

相关题目：CodeForces 76F. Tourist

观察：为使相同颜色的节点处在同一个子树中，则包含这些节点的最小子树的所有节点必然会被划分在同一部分。

因此，在随意选择一个节点作为树的根节点后，每种颜色的所有节点的LCA（最近公共祖先）必然也与这些节点在同一部分。

同时，我们也得到了无解判定：如果某两种颜色的节点的最小子树具有相同部分，则必定无解。

在判断有解之后，我们可以把每种颜色对应的最小子树缩成一个节点，则问题就转化为：

【一个$n \leq 3\times 10^5$个节点的树，其中有$k$个节点是被标记的，问有多少种方法把树分成$k$部分，每部分包含恰好一个被标记的节点。】

Step 2. 动态规划

我们在缩点之后，只需要解决转化后的问题。

设$f[x][s]$表示以$x$为根的子树有多少种划分方式，使得$x$所在的部分 【未包含$s=0$ / 包含$s=1$】 一个被标记的节点。于是答案为$f[r][1]$，其中$r$是根节点。

1. 若$x$未被标记，则

1.1. 若$x$所在部分未包含被标记的节点，则对每个$x$的儿子节点$y$，若$y$所在部分包含了被标记的节点，则必然不与$x$在同一部分；若$y$所在部分未包含被标记节点，则必然与$x$在同一部分，因此有$f[y][0]+f[y][1]$种可能。由乘法原理，有

$$ f[x][0] = \prod_{y \in \text{son}(x)} (f[y][0]+f[y][1]). $$

1.2. 若$x$所在部分包含被标记的节点，则枚举$x$的儿子节点$y$，其所在部分包含被标记节点，有$f[y][1]$种可能；对其他儿子节点$z \neq y$，若$z$所在部分包含了被标记的节点，则必然不与$x$在同一部分；若$z$所在部分未包含被标记节点，则必然与$x$在同一部分，因此有$f[z][0]+f[z][1]$种可能。由乘法原理和加法原理，有

$$ f[x][1] = \sum_{y \in \text{son}(x)} f[y][1] \prod_{y \neq z \in \text{son}(x)} (f[z][0]+f[z][1]). $$

2. 若$x$被标记，则

2.1. $x$所在部分不可能未包含被标记节点，即

$$ f[x][0] = 0, $$

2.2. 若$x$所在部分包含被标记的节点，则对每个$x$的儿子节点$y$，若$y$所在部分包含了被标记的节点，则必然不与$x$在同一部分；若$y$所在部分未包含被标记节点，则必然与$x$在同一部分，因此有$f[y][0]+f[y][1]$种可能。（这与1.1.的讨论相同）由乘法原理，有

$$ f[x][1] = \prod_{y \in \text{son}(x)} (f[y][0]+f[y][1]). $$

总时间复杂度为$O(n)$。