自分の解法において興味深い知見を得られたので復習も兼ねてメモを残します。

考察部分は省略します。自分の方針においては、整理すると以下の問題を解ければよいことになりました。

正整数 \(N, M, K\) が与えられる。(\(K \leq M\))

\[f(a) = \sum_{i=0}^{K-1} \binom{i}{a} \binom{M-i}{a} (K-i)\]

とする。\(f(0), f(1), \dots, f(N-1)\) を列挙せよ。

\(\sum_i\) の範囲に \(K\) の制約があることがポイントです。(範囲の制限が無ければ \(f(a)\) はヴァンデルモンドの恒等式の要領で簡単な二項係数になります。)

一見するとこの式は高速に計算できないようですが、手段を尽くすと \(\mathrm{O}(N \log^2 N + M)\) 程度で計算できます。

\[ \begin{aligned} f(a) &= \sum_{i=0}^{K-1} [x^a y^a] (1+x)^i (1+y)^{M-i} (K-i) \\ &= [x^a y^a] \sum_{i=0}^{K-1} (1+x)^i (1+y)^{M-i} (K-i) \end{aligned} \]

です。右辺の 2 変数母関数の部分に注目すると、

\[ \begin{aligned} &\sum_{i=0}^{K-1} (1+x)^i (1+y)^{M-i} (K-i) \\ &= (1+y)^M \sum_{i=0}^{K-1} \left(\frac{1+x}{1+y}\right)^i (K-i) \end{aligned} \]

と変形出来て、\(\sum_i (ai+b) c^i\) という形をしていることから等比数列の和の公式の要領で閉形式が求まります。実際に計算すると上式は

\[(1+y)^{M+1-K}\frac{(1+x)^{K+1}-(K+1)(1+x)(1+y)^K+K(1+y)^{K+1}}{(x-y)^2}\]

となります。この式は \(\frac{(1+x)^A (1+y)^B}{(x-y)^2}\) の和の形をしているため、結局

\[[x^a y^a] \frac{(1+x)^A (1+y)^B}{(x-y)^2}\]

を \(a\) について列挙する問題に帰着されます。ただし分子部分は \((x-y)^2\) で約せず分母が残るので、\((x-y)^{-2}\) のローラン展開(負冪を認める級数展開)においてどちらの変数を主変数、すなわち負冪を許さない変数として展開するかを指定する必要があり、今回は \(y\) を主変数とする展開を採用します。つまり、

\[\frac{1}{(x-y)^2} = x^{-2} \frac{1}{(1-y/x)^2} = \sum_{i\geq 0} (i+1) x^{-i-2} y^i\]

とします。より厳密に述べると、 \(\mathbb{F}((x^{-1}))\lbrack \lbrack y \rbrack \rbrack\) 、すなわち \(x^{-1}\) のローラン級数体を係数とする \(y\) の形式的冪級数全体からなる環上で展開します。

以降では

\[g_a = [x^a y^a] \frac{(1+x)^A (1+y)^B}{(x-y)^2}\]

とした時に \(g_0, g_1, \dots, g_{N-1}\) を列挙する問題を考えます。

こうした問題では \(x = z, y = u/z\) を代入することで定数項の取り出しに帰着させる手法が知られています。実際に代入すると

\[ \begin{aligned} g_a &= [z^0 u^a] \frac{(1+z)^A (1+u/z)^B}{(z-u/z)^2} \\ &= [z^0 u^a] (1+z)^A (1+u/z)^B \sum_{k\geq 0} (k+1) u^k z^{-2k-2} \\ &= [z^0 u^a] (1+z)^A \sum_{k \geq 0} \sum_{d \geq 0} (k+1) \binom{B}{d} u^{k+d} z^{-2k-2-d} \end{aligned} \]

ここで \([u^a]\) を取るには \(k+d=a\) が必要なので \(k\) を消去することができます。\(k \geq 0 \iff d \leq a\) という条件が残ることを踏まえて、

\[ \begin{aligned} &= \sum_{d=0}^{a} (a-d+1) \binom{B}{d} \left([z^0] (1+z)^A z^{-2a-2+d} \right) \\ &= \sum_{d=0}^{a} (a-d+1) \binom{B}{d} \binom{A}{2a+2-d} \end{aligned} \]

と変形することが出来ました。この式は

\[U(a) = \sum_{d=0}^a S(d) T(2a-d)\]

型の畳み込みが出来れば計算可能で、\(0 \leq d \leq a \iff 0 \leq d \leq 2a-d\) であることを踏まえると \(d=i,2a-d=j\) と置換すれば

\[U(a) = \sum_{i+j=2a, 0 \leq i \leq j} S(i) T(j)\]

型の畳み込みができればよいことになり、この式は分割統治 FFT で計算可能です。以上より全てを \(\mathrm{O}(N \log^2 N + M)\) 程度で計算することが出来ました。