edu118 F (容斥原理+NTT)

F. Tree Coloring

为 \(n\) 个结点分配 \(n\) 种颜色，求给定的 \(n-1\) 条限制均不违反的方案数。

设 \(f(k)\)：钦定其中的 \(k\) 条限制必须违反，钦定的总方案数。

对于每条违反的限制，可以看作是两个颜色的合并（即在树中，颜色 \(c\) 必须是 颜色 \(c - 1\) 的父亲，想象成 \(1 \backsim n\) 的一条链与并查集合并过程，每次合并等价于少了一个结点）。那么对于违反了钦定好的 \(k\) 条限制的情况，方案数为 \((n - k)!\)，证明：

• 对于每条限制，只需要考虑其上方的点填什么颜色 \(\rightarrow\) 先选择这 \(k\) 个点的颜色 \(\rightarrow\) 要求这 \(k\) 个颜色两两不相邻，且不能选第一个颜色（否则就选不了下方的点了） \(\rightarrow\) \(C_{n-k}^{k}\)（利用下面的插空法推导）
  
• 然后，有限制的 \(k\) 对点 与 无限制的 \(n-2k\) 个点分别可以任意排布 \(\rightarrow\) \((k!) * (n-2k)!\)
• 二者乘法原理 \(\rightarrow\) \(C_{n-k}^{k} * (k!) * (n-2k)! = (n-k)!\)

根据容斥原理，答案可以写作：

\[\sum_{i=0}^{n-1} (-1)^{i}*f_{i}*(n-i)! \]

问题在于如何求 \(f_{i}\)。

若所有限制是相互独立的，那么显然 \(f_{i} = C_{n}^{i}\)；而每条限制的违反需要在原树里的某条树边上体现，而对于同一个结点向下引出的两条边，它们一定不会对应同时违反的两条限制。故这些限制并不是相互独立的，需要用其他方式来计算 \(f_{i}\)。

\(f_{i}\) 等价于：在原树中选出 \(i\) 条边，满足对于树中的任意结点，向下引出的所有边中至多只能存在一条选中的边，方案数。

树是有根树，设结点 \(i\) 的孩子数量为 \(c_{i}\)，很明显 \(f_{i}\) 与 \(c_{i}\) 有关联。定义 \(dp_{i,j}\)：考虑结点 \(1 \backsim i\)，恰好选 \(j\) 条边，并满足上述条件，方案数。

状态转移：

\[dp_{i,j} = dp_{i-1,j-1}*c_{i} + dp_{i-1,j} \]

直接暴力转移复杂度是 \(O(n^{2})\) 的，考虑如何优化。

考虑利用 多项式 + NTT 优化。具体地，将 \(dp\) 的第二维 \(j\) 看作多项式幂次，将每一项的系数看作 \(dp\) 值。我们会发现一个神奇的事情：为每个结点 \(i\) 分配一阶多项式 \((1+c_{i}x)\)，那么 \(\prod_{i=1}^{n} (1+c_{i}x)\) 中 \(x^{j}\) 的系数就表示 \(dp_{n,j}\)。具体证明可以利用数学归纳法，结合上述定义分析。

对于多项式乘法，可以利用 NTT 快速计算。特别地，对于 \(n\) 个多项式作卷积，当它们的总项数很少时 (\(O(n)\))，可以利用多项式的启发式合并，将复杂度优化到 \(O(n \log^{2} n)\)。具体实现见 code。

code