「十二省联考 2019」希望（组合计数 + 动态规划 + 长链剖分）

题目大意

我们所可以自慰的，想来想去，也还是所谓对于将来的希望。
希望是附丽于存在的，有存在，便有希望，有希望，便是光明。

「十二省联考 2019」希望（LOJ 3053）

给定一棵 $n$ 个结点的树以及两个数 $L, k$。对于树上的一个连通快 $S$，定义 $R(S)$ 为在 $S$ 中且离 $S$ 中的每个点距离都不超过 $L$ 的点集。问任选 $k$ 个联通快 $S_1, S_2, \cdots, S_k$，使得 $R(S_1), R(S_2), \cdots, R(S_n)$ 有交的方案数 $\bmod 998244353$ 的结果。

数据范围：$n, L \le 10^6$。

思路分析

第一部分：初步转化

容易发现 $R(S)$ 也是一个联通块。联通块有如下的性质：

• 联通块和联通块的交还是联通块
• 非空联通块的点数等于边数 $+1$

于是答案就可以转化为：

$$\sum_{i=1}^{n} f(i)^k - \sum_{i=2}^{n} g(i)^k$$

其中 $f(i)$ 表示有多少个 $S$ 满足 $i \in R(S)$，$g(i)$ 表示有多少个 $S$ 满足 $i, \text{fa}(i) \in R(S)$。

第二部分：dp 建模

考虑 DP。设计如下状态：

• $\text{dp}(i, k)$ 满足在 $i$ 子树内有多少个包含 $i$ 的联通块，满足其中的所有点到 $i$ 的距离不超过 $k$
• $\text{up}(i, k)$ 表示在 $i$ 子树外（包含 $i$）有多少个包含 $i$ 的联通块，满足其中的所有点到 $i$ 的距离不超过 $k$。

发现 $f(i) = \text{dp}(i, L) \times \text{up}(i, L)$，$g(i) = \text{dp}(i, L - 1) \times (\text{up}(i, L) - 1)$。可以通过计算 $\text{dp}, \text{up}$ 得到 $f, g$。

列出状态转移方程：

• $\text{dp}(i, k) = \prod_{j \in \text{ch}(i)} \text{dp}(j, k - 1) + 1$
• $\text{up}(i, k) = \text{up}(\text{fa}(i, k - 1)) \times \prod_{j \in \text{ch}(\text{fa}(i)), j \neq i} \text{dp}(j, k - 2) + 1$

于是我们得到了 $O(nL)$ 的做法，期望得分 $36$ 分。

至此难度不大于 NOIP 提高组。

第三部分：长链剖分

最难部分。

这个 DP 是和深度有关的，容易想到长链剖分优化。

先考虑优化 DP 数组的转移。我们设当前结点为 $u$，当前儿子结点为 $v$。我们需要把短儿子依次合并到长儿子上面。

但是有一个问题：对于深度大于短儿子深度的部分，我们要将其整体乘上 $\text{dp}(v, d(v))$，其中 $d(i)$ 表示点 $i$ 往下延伸的最长长度。我们发现要乘上一个数的后缀长度很大，是 $O(d(u) - d(v))$ 级别的，但是长链剖分要求合并的复杂度为 $O(d(v))$ 才能够做到线性。

于是我们考虑维护整体乘标记。把后缀乘法操作转化为先整体乘一个数，再将对应的前缀乘这个数的逆元。我们还要维护整体加标记，因为每次做完一个结点后，我们要将它的 DP 值 $+1$。此外，为了方便转移，也要维护整体乘标记的逆元。

乍一看算法好像是正确的，其实不然。注意到 $\text{dp}(v, d(v))$ 可能不存在逆元，这样后缀乘法操作就变成了后缀变 $0$ 操作。所以还需维护一个后缀变 $x$ 标记，其中 $x$ 与乘法标记相乘再加上加法标记等于 $0$。

再考虑优化 $\text{up}$ 数组的转移。考虑先转移短儿子，再转移长儿子。

我们发现短儿子所用到的 $\text{up}$ 值的深度范围是 $d(v)$ 级别的，所以可以直接转移。

而长儿子可以打标记 + 将所有短儿子转移过来，复杂度为所有短儿子 $d$ 之和，也是均摊线性的。

现在主要的问题就是如何维护某个点的所有 “兄弟” 的 DP 值之积。发现它一定是一个前缀积和一个后缀积相乘。我们将 $u$ 的所有儿子倒过来考虑，前缀积可以通过在之前计算 $dp$ 值的时候将有更改的位置预先存起来，后缀积则可以边扫边更新。

此时离正解已经很接近了，可是求一个数逆元的复杂度还是带有 $\log$。所以时间复杂度为 $O(n \log n)$。

第四部分：离线求逆元

要做到严格线性，就要将求逆元的过程优化。

发现需要求的逆元只有 $h(i)$，其中 $h(i)$ 表示 $i$ 的子树内有多少个包含 $i$ 的联通块。这显然是可以预先线性地 DP 出来的。所以问题就变成了给定 $a_1, a_2, \cdots, a_n$，求 $a_1^{-1}, a_2^{-1}, \cdots, a_n^{-1} \mod 998244353$。

我们令 $A(i) = \prod_{j = 1}^{i} a_i$，可以先正着做一次前缀积求解 $A(i)$，然后通过 $A(n)$ 得出 $A(n)^{-1}$，最后反着做一次后缀积来求解 $A(i)^{-1}$。又有 $a_i = A(i)^{-1} \times A(i - 1)$，于是我们就可以线性求出逆元。

至此题目已经完美解决，时间复杂度 $O(n)$，期望得分 $100$ 分。

代码实现

#include <cstdio>