快速计算分割函数

· · 算法·理论

核心思路

  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) = \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)

我们需要计算 A(x) = \sum_{k=-\infty}^{\infty} (-1)^k x^{g_k} \pmod{x^{N+1}}。

我们只需要计算所有 g_k \le N 的项。 由于 g_k \approx \frac{3k^2}{2}， k 的取值范围大致在 O(\sqrt{N}) 级别。

  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 相同）。

复杂度：这个循环最多执行 O(\sqrt{N}) 次。所以构造 A(x) 的时间复杂度仅为 O(\sqrt{N})。

步骤 2：计算多项式求逆

我们现在有了 A(x) (一个 N+1 项的多项式，尽管它很稀疏)。 我们的目标是计算 P(x) 使得：

A(x) P(x) \equiv 1 \pmod{x^{N+1}}

这可以通过基于牛顿迭代法 的标准多项式求逆算法实现，从而可以在 O(N \log N) 的时间内计算出 A(x) 的逆 P(x)。

多项式求逆简介： 假设我们已经求得 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：获得答案

执行完多项式求逆后，得到的 P(x) = \sum_{n=0}^{N} p(n)x^n。 数组 P 中的第 n 项 P[n] 即为分割函数 p(n) 的值。

总时间复杂度：O(\sqrt{N} + N \log N) = O(N \log N)。

附：另一种 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}}。