F2. Tree Cutting (Hard Version)

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

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

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

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

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

Step 2. 动态规划

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

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

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̸=y$，若$z$所在部分包含了被标记的节点，则必然不与$x$在同一部分；若$z$所在部分未包含被标记节点，则必然与$x$在同一部分，因此有$f[z][0]+f[z][1]$种可能。由乘法原理和加法原理，有

$$f[x][1]=\sum _{y\in \text{son}(x)}f[y][1]\sum _{y̸=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)$。

原文：https://www.cnblogs.com/TinyWong/p/10422194.html