快速计算分割函数

核心思路

1. 找到分割函数的生成函数：分割函数 $p(n)$（将 $n$ 拆分为若干正整数之和的方案数）的生成函数 $P(x)$ 是： $P(x) = \sum_{n=0}^{\infty} p(n)x^n = \prod_{k=1}^{\infty} \frac{1}{1-x^k}$ 我们的目标是计算这个 $P(x)$ 在 $\pmod{x^{N+1}}$ 意义下的系数，即 $p(0), p(1), \dots, p(N)$。
2. 利用多项式求逆： $P(x)$ 可以看作是另一个多项式 $A(x)$ 的倒数： $P(x) = \frac{1}{A(x)} \quad \text{其中} \quad A(x) = \prod_{k=1}^{\infty} (1-x^k)$ 如果我们能快速计算出 $A(x) \pmod{x^{N+1}}$，我们就可以用一次 $O(N \log N)$ 的多项式求逆来得到 $P(x) \pmod{x^{N+1}}$。
3. 使用欧拉五边形数定理： 直接计算 $A(x) = (1-x)(1-x^2)\dots(1-x^N) \pmod{x^{N+1}}$ 需要 $O(N^2)$ 或更差的复杂度，这太慢了。 幸运的是，欧拉发现了一个神奇的定理——五边形数定理： $\prod_{k=1}^{\infty} (1-x^k) = \sum_{k=-\infty}^{\infty} (-1)^k x^{g_k} = 1 + \sum_{k=1}^{\infty} (-1)^k \left( x^{g_k} + x^{g_{-k}} \right)$ 其中 $g_k = \frac{k(3k-1)}{2}$ 是广义五边形数。
   展开来看 $A(x)$：
   • $k=0$: $g_0 = 0$, $A(x)$ 有 $+1 \cdot x^0 = 1$
   • $k=1$: $g_1 = 1$, $g_{-1} = 2$, $A(x)$ 有 $-1 \cdot x^1 - 1 \cdot x^2 = -x - x^2$
   • $k=2$: $g_2 = 5$, $g_{-2} = 7$, $A(x)$ 有 $+1 \cdot x^5 + 1 \cdot x^7 = +x^5 + x^7$
   • $k=3$: $g_3 = 12$, $g_{-3} = 15$, $A(x)$ 有 $-1 \cdot x^{12} - 1 \cdot x^{15}$
   • ...
   所以： $A(x) = 1 - x - x^2 + x^5 + x^7 - x^{12} - x^{15} + \dots$
   这是一个稀疏多项式！

$O(N \log N)$ 算法步骤

步骤 1：构造稀疏多项式 $A(x)$

1. 初始化一个全零的多项式（数组） $A$ ，长度为 $N+1$。
2. 设置 $A[0] = 1$。
3. 从 $k = 1$ 开始循环： a. 计算五边形数 $g_k = \frac{k(3k-1)}{2}$。 b. 如果 $g_k > N$，跳出循环。 c. 根据 $k$ 的奇偶性设置 $A[g_k]$： * 如果 $k$ 是奇数， $A[g_k] = -1$。 * 如果 $k$ 是偶数， $A[g_k] = +1$。 d. 计算 $g_{-k} = \frac{k(3k+1)}{2}$。 e. 如果 $g_{-k} > N$，跳出循环（理论上可以只跳出这一半）。 f. 同样设置 $A[g_{-k}]$（符号与 $g_k$ 相同）。

步骤 2：计算多项式求逆

多项式求逆简介： 假设我们已经求得 $P_0(x)$ 满足 $A(x)P_0(x) \equiv 1 \pmod{x^k}$。 我们想求 $P(x)$ 满足 $A(x)P(x) \equiv 1 \pmod{x^{2k}}$。

迭代公式为： $P(x) \equiv P_0(x) \left( 2 - A(x)P_0(x) \right) \pmod{x^{2k}}$

从 $P(x) \equiv 1 \pmod{x^1}$ (即 $p(0)=1$) 开始，我们通过 $\log N$ 次迭代（每次将精度加倍：$1 \to 2 \to 4 \to \dots \to N'$），每次迭代涉及两次多项式乘法（和一次减法），总复杂度为 $O(N \log N)$。

步骤 3：获得答案

附：另一种 $O(N \log N)$ 方法 (多项式 Exp)

1. $P(x) = \prod_{k=1}^{\infty} \frac{1}{1-x^k}$
2. 取对数：$\ln P(x) = \sum_{k=1}^{\infty} -\ln(1-x^k) = \sum_{k=1}^{\infty} \sum_{j=1}^{\infty} \frac{(x^k)^j}{j} = \sum_{k=1}^{\infty} \sum_{j=1}^{\infty} \frac{x^{kj}}{j}$
3. 交换求和次序，令 $n = kj$： $\ln P(x) = \sum_{n=1}^{\infty} x^n \left( \sum_{j|n} \frac{1}{j} \right) = \sum_{n=1}^{\infty} x^n \left( \frac{1}{n} \sum_{j|n} \frac{n}{j} \right) = \sum_{n=1}^{\infty} \frac{\sigma_1(n)}{n} x^n$ （其中 $\sigma_1(n)$ 是 $n$ 的约数和）
4. 令 $Q(x) = \ln P(x)$。
5. 算法： 使用 $O(N \log N)$ 的多项式指数函数 计算 $P(x) = \exp(Q(x)) \pmod{x^{N+1}}$。