F2. Tree Cutting (Hard Version)

觀察：爲使相同顏色的節點處在同一個子樹中，則包含這些節點的最小子樹的所有節點必然會被劃分在同一部分。

因此，在隨意選擇一個節點作爲樹的根節點後，每種顏色的所有節點的LCA（最近公共祖先）必然也與這些節點在同一部分。

同時，我們也得到了無解判定：如果某兩種顏色的節點的最小子樹具有相同部分，則必定無解。

在判斷有解之後，我們可以把每種顏色對應的最小子樹縮成一個節點，則問題就轉化爲：

【一個$n \leq 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 \neq 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 \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)$。

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