問題の言い換え

まず、\(F(1), F(2), \dots, F(M)\) をいきなり求めるのは少し難しいので、次のように \(a\) を置いてみましょう。

• \(a_n :=\) (\(C(P) = n\) である \(P\) の個数)

すると、この問題は次の 2 つのパートに分解できます。

• \(a_0, a_1, \dots, a_{N-1}\) を計算する
• \(a\) から \(F\) を求める

まずは \(a\) を計算します。

包除原理

\(a\) を陽に計算するのはまだ難しいです。そこで、次のように \(p, q\) を置きます。

• \(p_n :=\) (隣り合うボールの色が同じである場所が \(n\) ヵ所あるようなボールの列を構成する \(P\) の個数)
• \(q_n :=\) (少なくとも特定の \(n\) ヵ所について、隣り合うボールの色が同じであるようなボールの列を構成する \(P\) の個数)

明らかに \(a_n = p_{N-1-n}\) なので \(p\) が求められれば良いです。
\(p\) と \(q\) の間には次の関係が成り立ちます。

\[q_{n} = \sum_{k=n}^{N-1} \binom{k}{n} p_k\]

包除原理より、ここから \(p\) に関する式が得られます。

\[p_n = \sum_{k=n}^{N-1} (-1)^{k-n}\binom{k}{n} q_k \]

\(q\) から \(p\) は畳み込みで \(\mathrm{O}(N \log N)\) で求められます。よって \(q\) を高速に計算できればよいです。

指数型母関数を利用した \(q\) の計算

\(q_0, q_1, \dots, q_{N-1}\) を求めます。

色 \(c\) のボールが \(m_c\) 個あるとします。隣り合う可能性のある箇所は高々 \(m_c-1\) ヵ所ですが、このうち少なくとも特定の \(k_c\) ヵ所が隣り合っているとしましょう。隣り合っている箇所の選び方は \(\binom{m_c-1}{k_c}\) 通りあり、これを \(a_{c, k}\) とおきます。これを利用して適切な考察を経ると、\(q_n\) は次の式で計算できるとわかります。

\[q_n = \sum_{k_1, k_2, \dots, k_N ,\sum_{c} k_c = n} \left(\binom{N-n}{m_1-k_1, m_2-k_2, \dots, m_N-k_N}\prod_{c} a_{c,k}\right)\]

ここで、式を簡単にするために

\[r_0=0, r_i = q_{N-i} (i \geq 1)\]

を満たす \(r_0, r_1, \dots, r_N\) を置くと

\[r_n = \sum_{k_1, k_2, \dots, k_N ,\sum_{c} k_c = n} \left(\binom{n}{k_1, k_2, \dots, k_N}\prod_{c} a_{c,m_c - k}\right)\]

と表せます。
この式を用いると(少し難しいですが) DP によって \(r\) を \(\mathrm{O}(N^2)\) ですべて計算できますが、さらに高速化できます。

\[g_c(x) = \sum_{k=0}^{m_c} \frac{a_{c,m_c - k}}{k!} x^k\]

とおくと、

\[r_n = n! \times [x^n] \prod_{c} g_c(x)\]

という風に \(r_n\) は \(\prod_{c} g_c(x)\) の \(n\) 次の係数として表せます。よって、\(\prod_{c} g_c(x)\) を畳み込み + マージテクで \(\mathrm{O}(N \log^2 N)\) で計算すれば \(r\) を列挙できて、ここから \(q\) も求まります。

この種の畳み込みの背景として 指数型母関数 という概念があります。数列 \((a_n)\) に対して \((a_n)\) の指数型母関数 \(A(x)\) とは次のような関数です。

\[A(x) = \sum_{n=0}^\infty \frac{a_n}{n!} x^n\]

指数型母関数の特徴として、指数型母関数同士の畳み込みによって得られる母関数が二項係数の重みつきの畳み込みになるという点があります。

\[ \begin{aligned} &A(x) B(x) \\ &= \left(\sum_{n=0}^\infty \frac{a_n}{n!} x^n\right)\left(\sum_{m=0}^\infty \frac{b_m}{m!} x^m\right) \\ &= \sum_{l=0}^\infty \left( \sum_{n+m=l} \frac{a_n b_m}{n! m!} \right) x^l\\ &= \sum_{l=0}^\infty \left( \sum_{n+m=l} \binom{l}{n} a_n b_m \right) \frac{x^l}{l!}\\ \end{aligned} \]

上記の \(r\) の求め方は、この式を多項係数に拡張したものです。
以上より \(q\) を求めることができて、ここから \(a\) も求められました。

\(F(n)\) の通常型母関数

\[F(n) = \sum_{k=0}^{N-1} k^n a_k\]

です。ここで、\(F(n)\) の通常型母関数 \(f(x)\) を次のように置きます。

\[f(x) = \sum_{n=0}^\infty F(n) x^n\]

すると、\(f(x)\) は \(a\) を用いて次のように表現できます。

\[ \begin{aligned} &f(x) \\ &=\sum_{n=0}^\infty F(n) x^n \\ &=\sum_{n=0}^\infty \left( \sum_{k=0}^{N-1} k^n a_k \right) x^n \\ &=\sum_{k=0}^{N-1} a_k \left( \sum_{n=0}^\infty (kx)^n\right) \\ &=\sum_{k=0}^{N-1} \left(\frac{a_k}{1-kx}\right) \end{aligned} \]

よって、

\[\sum_{k=0}^{N-1} \left(\frac{a_k}{1-kx}\right)\]

の \(1\) 次から \(M\) 次 までの係数が求まればよいです。上の式に求めた \(a\) の値を代入して、畳み込み + マージテクを利用して総和を計算することで、

\[f(x) = \frac{s(x)}{t(x)}\]

という式を \(\mathrm{O}(N \log^2 N)\) で得ることができます。(\(s(x),t(x)\) は\(\deg(s) \lt \deg(t)\) である \(N\) 次以下の多項式)

形式的冪級数の逆元

以上より、\(\frac{s(x)}{t(x)}\) の冪級数展開を \(M\) 次の項まで計算できれば良いです。ここで次のような問題を考えます。

\(t_0, t_1, \dots, t_M\) が与えられる。ここで\(t_i\) は \(t(x)\) の \(i\) 次の係数であり、 \(t_0 = 1\) である。
次の条件を満たす \(u_0, u_1, \dots, u_M\) を求めよ。

• \(u_0 = 1\)
• \(n \geq 1\) について \(\sum_{i=0}^n t_i u_{n-i} = 0\)

このとき、\(u_n\) は \(u_0, u_1, \dots, u_{n-1}\) を用いて

\[u_n = - \sum_{i=1}^{n} t_i u_{n-i}\]

と表せるので、上の条件を満たす \(u\) は存在します。このような \(u_0, u_1, \dots, u_M\) の母関数を \(u(x)\) とします。すると、

\[t(x)u(x) \equiv 1 \pmod{x^{M+1}}\]

および

\[f(x) \equiv s(x) u(x) \pmod{x^{M+1}}\]

が成り立つことが確認できます。よって、\(u(x)\) が求まれば \(f(x)\) の \(M\) 次以下の項を \(s(x)\) と \(u(x)\) の畳み込みで求めることができます。

このように、形式的冪級数 \(t(x)\) に対して

\[t(x)u(x) \equiv 1 \pmod{x^{M+1}}\]

を満たす \(u(x)\) を \(t(x)\) の(乗法的)逆元と呼びます。

\(u(x)\) を求めます。\(n\) 次の係数 \(u_n\) は先に書いた \(u_n\) に関する式を利用して分割統治 FFT で \(\mathrm{O}(M \log^2 M)\) で求まりますが、ダブリングを利用するとより高速に求めることができます。
多項式 \(u_k\) を \(u_k := u(x) \bmod x^k\) とおきます。\(u_k\) が分かっているとき位に \(u_{2k}\) を高速に求められればよいです。

\[u_k - u_{2k} \equiv 0 \pmod{x^k}\]

なので

\[(u_k - u_{2k})^2 \equiv u_k^2 - 2 u_k u_{2k} + u_{2k}^2 \equiv 0 \pmod{x^{2k}}\]

です。両辺に \(t\) を掛けて

\[t u_k^2 - 2 t u_k u_{2k} + t u_{2k}^2 \equiv 0 \pmod{x^{2k}}\]

を得ます。\(t u_{2k} \equiv 1 \pmod{x^{2k}}\) を利用して変形すると

\[u_{2k} \equiv 2 u_k - t u_k^2 \pmod{x^{2k}}\]

を得ます。この式を利用して \(u_1=1\) からスタートして \(u\) の長さを倍々にしていけば \(\mathrm{O}(M \log M)\) で \(u(x)\) の \(M\) 次以下の係数を求められます。

なお、これをより一般化した手法として多項式のニュートン法と呼ばれるアルゴリズムがあります。(以下の式に \(A(F) = F \times t(x) - 1\) を代入すると上で説明したアルゴリズムと一致します。)

多項式のニュートン法

\(A(F) = 0\) を満たす \(F\) を求めたい。(\([x^0]F\) はもともとわかっているとする。) \(\hat{F}(x) = F(x) \bmod {x^d}\) がわかっているときに

\[A(F) = A(\hat{F}) + A'(\hat{F})(F - \hat{F}) + \mathrm{O}((F - \hat{F})^2)\]

の両辺を \(\bmod{x^{2d}}\) を取ると \(\mathrm{O}((F - \hat{F})^2)\) が無視できるので

\[A(F) \equiv A(\hat{F}) + A'(\hat{F})(F - \hat{F}) \pmod{x^{2d}}\]

から \(F \mod x^{2d}\) を計算できる。

以上よりこの問題を \(\mathrm{O}(N \log^2 N + M \log M)\) で解くことができました。