形式冪級数への帰着

条件 \(\max(B) = m\) の代わりに、条件 \(\max(B) \leq m\) を考えて最後に差分を取ることにします。

\(\max(B) \leq m\) を満たす長さ \(n\) の良い数列の組の個数を \(C_{n,m}\) と置きます。 ただし便宜上、すべての \(m \geq 0\) について \(C_{0,m} = 1\) とします。

まず \(B\) を固定したときの答えを考えます。 各 \(b\) について \(B_i = b\) なる \(i\) の個数を \(n_b\)、\(B_i \geq b\) なる \(i\) の個数を \(S_b\) とします。 このとき、\(B_i = b\) なる \(i\) に対して選べる \(A_i\) は \(B_j \geq b\) なる \(j\) なので、\(S_b\) 個です。 どの \(j\) を選ぶ回数も素数回にはなってはいけないことに注意すると、\(P(X)\) を素数ではない非負整数についての指数型母関数、すなわち

\[P(X) = \sum_{k \geq 0:\text{素数ではない}} \dfrac{X^k}{k!}\]

と定めれば、\(B_i = b\) なる \(i\) に対する \(A_i\) の選び方は

\[\left[\dfrac{X^{n_b}}{{n_b}!}\right] P(X)^{S_b}\]

通りです。\((n_1, \ldots, n_m)\) を固定したときの \(B\) の個数は

\[\dfrac{n!}{\prod\limits_{b=1}^m {n_b}!}\]

なので、

\[\dfrac{C_{n,m}}{n!} = \sum_{n_1 + \cdots + n_m = n}\prod_{b=1}^m [X^{n_b}] P(X)^{S_b}\]

となります。

\(\dfrac{C_{s,m}}{s!}\) から \(\dfrac{C_{t,m+1}}{t!}\) への遷移を考えると、係数は \([X^{t-s}] P(X)^t\) です。

\(Q(X)\) を \(\dfrac{X}{P(X)}\) の逆関数とします。Lagrange 反転公式により、

\[[X^{t-s}] P(X)^t = [X^t]\dfrac{XQ'(X)}{Q(X)} Q(X)^s\]

となるので、\(F_m(X) = \displaystyle \sum_{n = 0}^\infty \dfrac{C_{n,m}}{n!}X^n\) とすると、

\[F_{m+1}(X) = \dfrac{XQ'(X)}{Q(X)} F_m(Q(X))\]

となります。 ここから、帰納的に \(F_{m}(X) = \dfrac{X(Q^{\langle m \rangle})'(X)}{Q^{\langle m \rangle}(X)}\) であることがわかります。ここで、\(Q^{\langle m \rangle}(X)\) は \(Q(X)\) の \(m\) 階合成冪です。 実際、\(m = 0\) では明らかであり、\(m\) について正しければ、

\[\begin{aligned} F_{m+1}(X) =& \dfrac{XQ'(X)}{Q(X)} F_m(Q(X)) \\ =& \dfrac{XQ'(X)}{Q(X)} \dfrac{Q(X)(Q^{\langle m \rangle})'(Q(X))}{Q^{\langle m+1 \rangle}(X)} \\ =& \dfrac{X(Q^{\langle m+1 \rangle})'(X)}{Q^{\langle m+1 \rangle}(X)} \end{aligned}\]

より、\(m+1\) についても正しいです。

よって、\([X^N] \dfrac{X(Q^{\langle m \rangle})'(X)}{Q^{\langle m \rangle}(X)}\) を \(m = 1, \ldots, M\) について求める問題に帰着しました。

共役を用いた計算

\(Q = G \circ \dfrac{X}{1-X} \circ G^{-1}\) なる形式冪級数 \(G(X) = X + O(X^2)\) が存在します。 \(G\) の存在の証明や計算方法は後述します。 \(G\) が求まったときに元の問題の答えを求める方法を説明します。

\(\dfrac{X}{1-X}\) の \(m\) 階合成冪は \(\dfrac{X}{1-mX}\) なので、 \(Q^{\langle m \rangle} = G \circ \dfrac{X}{1-mX} \circ G^{-1}\) となり、

\[\dfrac{X(Q^{\langle m \rangle})'(X)}{Q^{\langle m \rangle}(X)} = \dfrac{X \left(G\left(\dfrac{G^{-1}(X)}{1-mG^{-1}(X)}\right)\right)'}{G\left(\dfrac{G^{-1}(X)}{1-mG^{-1}(X)}\right)}\]

となります。 ここで Lagrange 反転を用いることで、

\[[X^N]\dfrac{X \left(G\left(\dfrac{G^{-1}(X)}{1-mG^{-1}(X)}\right)\right)'}{G\left(\dfrac{G^{-1}(X)}{1-mG^{-1}(X)}\right)} = [X^N] A(X)^N B\left(\dfrac{X}{1-mX}\right) \dfrac{1}{1-mX}\]

となります。 ここで \(A(X) = \dfrac{X}{G(X)}, \ B(X) = \dfrac{XG'(X)}{G(X)}\) です。 \(A(X)^N = \displaystyle\sum_k a_k X^k, \ B(X) = \displaystyle\sum_k b_k X^k\) とおくと、

\[[X^l] B\left(\dfrac{X}{1-mX}\right)\dfrac{1}{1-mX} = \sum_k b_k [X^l]\dfrac{X^k}{(1-mX)^{k+1}} = \sum_k b_k \binom{l}{k} m^{l-k}\]

より、

\[[X^N] A(X)^N B\left(\dfrac{X}{1-mX}\right)\dfrac{1}{1-mX} = \sum_l\sum_k a_{N-l} b_k \binom{l}{k} m^{l-k}\]

となり、\(m^k\) の係数は \(\displaystyle\sum_{i-j=k} \binom{i}{j} a_{N-i} b_j\) となります。 適当な列の畳み込みにより、\(k = 0, \ldots, N\) についてこの値を求めることができ、 \([X^N] \dfrac{X(Q^{\langle m \rangle})'(X)}{Q^{\langle m \rangle}(X)}\) を \(m\) の多項式として表したときの係数を求めることができます。 最後に多項式の多点評価を行うことで、\(m = 1, \ldots, M\) に対する値が求まります。 （\(N \leq M\) のときは \(m = 1, \ldots, N+1\) についての値を求めて、残りの値を適当な畳み込みで計算することで計算量を改善できます。） 以上により、\(G\) の計算を除いて元の問題を時間計算量 \(O(N \log^2 N + M \log M)\) で解くことができます。

\(G\) の計算

以下の一般的な問題に帰着します。

• 形式冪級数 \(F(X) = X + X^2 + X^3 + O(X^4)\) が与えられるので、\(G^{-1} \circ F \circ G(X) \equiv \dfrac{X}{1-X} \pmod{X^N}\) なる形式冪級数 \(G\) を求めよ。

この問題を時間計算量 \(O(N \log^2 N)\) で解く方法を説明します。

Newton法でこの問題を解く方法を考えます。 \(n \geq 4\) に対して、\(G_0\) が

\[G_0^{-1} \circ F \circ G_0(X) \equiv \dfrac{X}{1-X} \pmod{X^n}\]

を満たしていたとして、

\[G^{-1} \circ F \circ G(X) \equiv \dfrac{X}{1-X} \pmod{X^{2n-3}}\]

を満たす \(G\) を \(O(n\log^2 n)\) で見つけられればよいです。 \(F\) を \(G_0^{-1} \circ F \circ G_0\) で、\(G\) を \(G_0^{-1} \circ G\) で置き換えることで、\(F \equiv \dfrac{X}{1-X} \pmod{X^n}\) としてよいです。これらの変換は \(O(n\log^2 n)\) で計算できます。

\(G(X) = X + X^{n-1} g(X)\) とおきます。 このとき、\(G^{-1}(X) \equiv X - X^{n-1} g(X) \pmod{X^{2n-3}}\) が成り立ちます。 実際、

\[\begin{aligned} G(X - X^{n-1} g(X)) &= X - X^{n-1} g(X) + (X - X^{n-1} g(X))^{n-1} g(X - X^{n-1} g(X)) \\ &\equiv X - X^{n-1} g(X) + X^{n-1} g(X) \\ &\equiv X \pmod{X^{2n-3}} \end{aligned}\]

となります。 また、\(F = \dfrac{X}{1-X} + X^n H(X)\) とおきます。 \((G^{-1})'(X) \equiv 1 \pmod{X^{n-2}}\) に注意すると、

\[\begin{aligned} G^{-1} \circ F \circ G(X) &= G^{-1}\left(\dfrac{G(X)}{1-G(X)} + G(X)^n H(G(X))\right)\\ &\equiv G^{-1}\left(\dfrac{G(X)}{1-G(X)}\right) + G(X)^n H(G(X)) (G^{-1})'\left(\dfrac{G(X)}{1-G(X)}\right)\\ &\equiv \dfrac{G(X)}{1-G(X)} - \left(\dfrac{G(X)}{1-G(X)} \right)^{n-1} g\left(\dfrac{G(X)}{1-G(X)} \right) + X^n H(X)\\ &\equiv \dfrac{X}{1-X} + \dfrac{X^{n-1}g(X)}{(1-X)^2} - \left(\dfrac{X}{1-X} \right)^{n-1} g\left(\dfrac{X}{1-X} \right) + X^n H(X) \pmod{X^{2n-3}} \end{aligned}\]

となります。 したがって、\(G^{-1} \circ F \circ G(X) \equiv \dfrac{X}{1-X} \pmod{X^{2n-3}}\) となる条件は、

\[\dfrac{X^{n-1}g(X)}{(1-X)^2} - \left(\dfrac{X}{1-X} \right)^{n-1} g\left(\dfrac{X}{1-X} \right) + X^n H(X) \equiv 0 \pmod{X^{2n-3}}\]

となります。両辺を \(\left(\dfrac{X}{1-X}\right)^{n-1}\) で割ることで、

\[(1-X)^{n-3} g(X) - g\left(\dfrac{X}{1-X}\right) + X(1-X)^{n-1} H(X) \equiv 0 \pmod{X^{n-2}}\]

を得ます。 \(g(X) = \displaystyle \sum_{i \geq 0} c_i X^i\) とおき、左辺の各項を展開して \(X\) の係数を比較すると、低次の \(c_i\) から順に決定することができます。 具体的には、左辺の \(X^i\) の係数は、\(g\) のより低次の項から適当な畳み込みを用いることで寄与が計算できます。 よって、分割統治FFT（オンライン畳み込み）を用いることで、\(O(n \log^2 n)\) で \(g\) を求めることができます。

以上により、\(O(N \log^2 N)\) で \(G(X)\) の先頭 \(N\) 項を計算できます。