CF1118F2 Tree Cutting (Hard Version)

题意

给定一个有 $n$ 个节点的树，节点可能有颜色，共 $k$ 种颜色，编号 $1...k$，保证每种颜色都出现. 有的点没有颜色，用 $0$ 表示. 将其划分为 $k$ 个联通块，是每个联通块中有且仅有一种颜色，颜色为 $0$ 的节点可以在任意联通块中. 求有多少中划分的方案

$2\le k\le n\le 3\times 10^5$，答案对 $998244353$ 取模

题解

显然有一些点的染色是确定的，一些点可以染上多种颜色，其中所有颜色为 $c$ 的点与它们的LCA的路径上的点确定染上颜色 $c$
先将这些颜色确定的点染色

接下来考虑树形DP
令 $f_{u,0}$ 表示在以 $u$ 为根的子树中，最上方联通块不包含已染色点的方案数，$f_{u,1}$ 表示在以 $u$ 为根的子树中，最上方联通块包含已染色点的方案数，$c_u$ 表示点 $u$ 的颜色

分两种情况讨论转移：

• $c_u\not=0$ 显然 $f_{u, 0}=0$
  对于 $f_{u, 1}$，再分情况讨论：

  • $v\in\text{son}_u,c_v\not=0$
    能够确定边 $(u, v)$ 是否删去，对 $f_{u, 1}$ 贡献为 $f_{v, 1}=f_{v, 0}+f_{v, 1}$
  • $v\in\text{son}_u,c_v=0$
    若边 $(u, v)$ 删去，$v$ 所在联通块必须有色，对 $f_{u, 1}$ 贡献为 $f_{v, 1}$
    若边 $(u, v)$ 保留，$v$ 所在联通块必须无色，对 $f_{u, 1}$ 贡献为 $f_{v, 0}$
    总贡献为 $f_{v, 0}+f_{v, 1}$ 综上，$f_{u, 1}=\prod\limits_{v\in\text{son}_u}f_{v, 0}+f_{v, 1}$

• $c_u=0$
  类似的，$f_{u, 0}=\prod\limits_{v\in\text{son}_u}f_{v, 0}+f_{v, 1}$
  对于 $f_{u, 1}$，枚举 $u$ 继承哪一个儿子节点的颜色，即 $f_{u, 1}=\sum\limits_{v_1\in\text{son}_u}f_{v_1, 1}\times\prod\limits_{v_2\in\text{son}_u,v_2\not=v_1}f_{v_2, 0}+f_{v_2, 1}$
  对这条式子维护 $f_{v, 0}+f_{v, 1}$ 的前缀积和后缀积即可

总的时间复杂度 $\Theta(n)$ 到 $\Theta(n\log n)$，取决于求LCA的算法

代码 codeforces submission 143330738