树状数组逆变换 题解

思路

把 $f$ 看成模 $998244353$ 意义下的线性变换。它的矩阵是下三角矩阵，且对角线全为一，所以逆变换唯一存在。

令 $F=I+N$，其中 $I$ 是单位变换，$N=F-I$ 是严格下三角部分。树状数组区间的包含关系使任意长度超过 $\lfloor\log_2 n\rfloor$ 的严格下降链不存在，因此 $N^j=0$（$j$ 足够大，且 $n\le2\cdot10^5$ 时取 $j=20$ 已足够）。

做法

子任务 1 每次用前缀和在线性时间内求出 $f^{-1}$，连续执行 $k$ 次，复杂度 $O(nk)$。

子任务 2 显式构造 $f^{-1}$ 的稠密矩阵，并用矩阵快速幂计算其 $k$ 次幂，复杂度 $O(n^3\log k)$。

满分算法使用负二项式展开：

$$F^{-k}=(I+N)^{-k}=\sum_{j\ge0}(-1)^j\binom{k+j-1}{j}N^j.$$

系数递推为

$$c_0=1,\qquad c_j=-c_{j-1}\frac{k+j-1}{j}.$$

对一个向量应用 $N$ 时，先用前缀和在线性时间计算 $F(v)$，再减去 $v$。依次求出 $b,Nb,N^2b,\ldots$ 并按系数累加即可。

正确性证明

$F$ 的对角线全为一，所以 $N=F-I$ 严格下三角且幂零。形式幂级数 $(1+z)^{-k}$ 的系数正是 $(-1)^j\binom{k+j-1}{j}$；把 $z$ 替换为幂零线性算子 $N$ 后级数有限终止，并满足 $(I+N)^kF^{-k}=I$。因此算法输出的 $a=F^{-k}b$ 必满足 $F^k a=b$。

复杂度

最多计算 $O(\log n)$ 个 $N$ 的幂向量，每次 $O(n)$，总时间复杂度 $O(n\log n)$，空间复杂度 $O(n)$。