トップ > 競技プログラミング > 分割数 p(n) を高速に求める

分割数 p(n) を高速に求める

競技プログラミング

競プロにおいて分割数は p(n, k) = n を k 個の非負整数の和で表す方法の数 (ただし, 順番は区別しない) であり, 計算量オーダーが \Theta(nk) である dp が(蟻本などに載っており)有名である

しかし, 一般に数論の世界で分割数というと p(n) = n を自然数の和で表す方法 であり, p(n) = p(n, n) である

これは同じ dp をすると \Theta(n ^ 2) かかってしまうが, n = k である場合に限りより高速な方法がある

まず, p(n) の母関数を考える

1 を何個使うか, 2 を何個使うか.... ということを考えると,

 \begin{align}
\displaystyle\sum _ {n  = 0} ^ {\infty} p(n) x ^ n &= (1 + x + x ^ 2 + \cdots )(1 + x ^ 2 + x ^ 4 + \cdots ) \cdots \cr
&= \prod _ {n  = 1} ^ {\infty} (1 + x ^n + x ^ {2n} + \cdots )
\end{align}

と書くことができる

 1 + x + x ^ 2 + \cdots = \dfrac{1}{1 - x}

と表せるため,

 \begin{align}
\displaystyle\sum _ {n  = 0} ^ {\infty} p(n) x ^ n &= \prod _ {n  = 1} ^ {\infty} \dfrac{1}{1 - x ^ n}
\end{align}

として, p(n) の母関数を求めることができた

できたが, 求めるのが簡単にはなっていない

ここで, 右辺の逆数である \prod _ {n = 1} ^ {\infty}(1 - x ^ n) について考えてみる

これは, オイラーの五角数定理 より

 \begin{align}
\prod _ {n  = 1} ^ {\infty}(1 - x ^ n) = \displaystyle\sum _ {n  = -\infty} ^ {\infty} (-1) ^ n x ^ {n(3n - 1)/2}
\end{align}

と書くことができる

凄まじく非自明だが, 証明も大変

ヤコビの三重積 から証明する方法, ヤング図形との対応から頑張る方法, Shanksによる天下り式な証明 などがあります 暇だったら書きます

0 \le n \le N について求められればいいので

 \begin{align}
\prod _ {n  = 0} ^ {N}(1 - x ^ n) = 1 + \displaystyle\sum _ {1 \le n \land n(3n - 1)/2 \le N} (-1) ^ n x ^ {n(3n - 1)/2} + \displaystyle\sum _ {1 \le n \land n(3n + 1)/2 \le N} (-1) ^ n x ^ {n(3n + 1)/2}
\end{align}

と書ける

これは簡単に求めることができて, NTT-friendly な mod でならニュートン法を用いて逆数を O(n\log n) で求めることができる

そうでない場合, さらに式変形をする

 \begin{align}
\displaystyle\sum _ {n  = 0} ^ {\infty} p(n) x ^ n &= \prod _ {n  = 1} ^ {\infty} \dfrac{1}{1 - x ^ n}
\end{align}

と

 \begin{align}
\prod _ {n  = 1} ^ {\infty}(1 - x ^ n) = \displaystyle\sum _ {n  = -\infty} ^ {\infty} (-1) ^ n x ^ {n(3n - 1)/2}
\end{align}

を掛け合わせて

 \begin{align}
\left( \displaystyle\sum _ {n  = 0} ^ {\infty} p(n) x ^ n \right) \left(  \displaystyle\sum _ {n  = -\infty} ^ {\infty} (-1) ^ n x ^ {n(3n - 1)/2} \right) = 1
\end{align}

と書ける

x ^ n の項の係数に注目することで

 \begin{align}
p(n) - p(n-1) - p(n-2) + p(n-5) \cdots = 0
\end{align}

という式が出てくる

したがって, 漸化式

 \begin{align}
p(0) &= 1 \cr
p(n) &= \displaystyle\sum _ {1 \le k \land k(3k - 1)/2 \le n} (-1) ^ {k-1} p\left(n - \dfrac{k(3k - 1)}{2}\right) + \displaystyle\sum _ {1 \le l \land l(3l + 1)/2 \le n} (-1) ^ {l-1} p\left(n - \dfrac{l(3l + 1)}{2}\right)
\end{align}

が得られる

この式を用いて愚直に求めていくことで, O(n \sqrt{n}) で求めることができる

Library Checker さんの Partition Function で verify できる