CF1118F2 Tree Cutting (Hard Version)

题意

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

，答案对 取模

题解

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

接下来考虑树形DP
令 表示在以 为根的子树中，最上方联通块不包含已染色点的方案数， 表示在以 为根的子树中，最上方联通块包含已染色点的方案数， 表示点 的颜色

分两种情况讨论转移：

• 显然
  对于 ，再分情况讨论：

  • 
    能够确定边 是否删去，对 贡献为
  • 
    若边 删去， 所在联通块必须有色，对 贡献为
    若边 保留， 所在联通块必须无色，对 贡献为
    总贡献为 综上，

• 
  类似的，
  对于 ，枚举 继承哪一个儿子节点的颜色，即
  对这条式子维护 的前缀积和后缀积即可

总的时间复杂度 到 ，取决于求LCA的算法

代码 codeforces submission 143330738