C - ABS Ball 解説 by shobonvip


形式的べき級数を用いて考察します。

\(i\) 番目の箱について、青のボールが赤のボールより \(k\) 個多いとすると、 \(|a_i - b_i| = k\) です。このとき、赤のボールの個数を \(s\) として、合計 \(2s+k\) 個のボールが使われています。

これを \(0\le s\)\(0 \le k\) のすべてについて考えて、さらに赤のボールが青のボールより多いという状況も同様に考えます。ボール \(1\) 個ごとに変数 \(x\) に対応させると、 \(i\) 番目の箱に対応する \(|a_i-b_i|\) の総和の形式的べき級数は

\[ 2(1+x^2+x^4+\cdots)(x+2x^2+3x^3+4x^4+\cdots) = \frac{2x}{(1-x^2)(1-x)^2}\]

となります。ここで、係数に掛けられた \(2\) は「青のボールが赤のボールより多い」と「赤のボールが青のボールより多い」という対称な \(2\) つのパターンを表しています。なお、赤のボールと青のボールの個数が同じ場合は、 \(|a_i - b_i| = 0\) となるため、考慮する必要がありません。

箱が \(M\) 個あり、求めるものは \(\prod_{i=1}^M |a_i-b_i|\) であるので、答えは

\[ \begin{aligned} [x^N] \prod_{i=1}^M \frac{2x}{(1-x^2)(1-x)^2} &= [x^N] \frac{2^M x^M}{(1-x^2)^M(1-x)^{2M}} \\ &= 2^M [x^{N-M}] \frac{1}{(1-x^2)^M(1-x)^{2M}} \end{aligned} \]

となります。いま、

\[ \begin{aligned} \frac{1}{(1-x^2)^M} &= \sum_{i\ge 0} \binom{M-1+i}{i} x^{2i}\\ \frac{1}{(1-x)^{2M}} &= \sum_{i\ge 0} \binom{2M-1+i}{i} x^{i} \end{aligned} \]

であるので(負の二項定理)、 \(i=0,1,\cdots, \lfloor (N-M)/2 \rfloor\) について、 \(j=N-M-2i\) として \(\binom{M-1+i}{i}\binom{2M-1+j}{j}\) を答えに加算すればよいです。つまり、答えは

\[ \sum_{i=0}^{\lfloor (N-M)/2 \rfloor} \binom{M-1+i}{i}\binom{M+N-1-2i}{N-M-2i} \]

となります。階乗とその逆数 (mod 998244353) の前計算を行うことで、この問題が \(O(N+M)\) で解けます。

投稿日時:
最終更新: