CF1118F2 Tree Cutting (Hard Version)

更好的阅读体验

题意

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

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

题解

我们称将一个点划归一个有且仅有颜色 \(c\) 的联通块为将这个点染上颜色 \(c\)

显然有一些点的染色是确定的，一些点可以染上多种颜色，其中所有颜色为 \(c\) 的点与它们的LCA的路径上的点确定染上颜色 \(c\)，证明显然

我们先按照如下流程将这些颜色确定的点染色：

1. 找出同一颜色 \(c\) 的点的LCA
2. 依次从每个点开始往父亲上跳，一边跳一边染色
3. 跳到一个点已经染色或已经到达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