公式解説より計算量が悪く、 \(O(N \sqrt{N \log N})\) の平方分割系の解法です。コンテスト中にG問題を解いたRating 2000以上の17人のうち平方分割系は2人いました。

問題は、公式解説のように

\[ \sum_{i=1}^{2N} C_i x^{P_i} (1+x)^{Q_i} \]

の展開を mod \(998244353\) で求めることに帰着されます。 ただし、数列 \(P, Q\) は

• \(0 \le P_i, Q_i\)
• \(P_i+Q_i \le N\)
• \(|P_{i+1} - P_i| \le 1\)
• \(|Q_{i+1} - Q_i| \le 1\)

を満たします。差分が高々 1 であることを利用して平方分割→FFTで解きます。

\(2N\) をいくつかのブロック \(B_1, \cdots, B_{\lceil 2N/K \rceil}\) に分けます。ブロックの大きさは最後を除きどれも \(K\) として、ブロック \(i\) には \(K\times (i-1)+1, \cdots, \min(2N, K\times i)\) の要素があるとします

それぞれのブロック \(B_i\) 内における \(P_j\) の最小値を \(P_{min}\)、 \(Q_j\) の最小値を \(Q_{min}\) として、\(B_i\) 内のすべての \(C_i x^{P_i} (1+x)^{Q_i}\) の和を、多項式 \(R_i(x)\) を用いて

\[ R_i(x) x^{P_{min}} (1+x)^{Q_{min}} \]

と表すことができます。ここで、 \(R_i(x)\) は

\[ R_i(x) = \sum_{i=K\times (i-1)+1}^{\min(2N, K\times i)} C_i x^{P_i - P_{min}} (1+x)^{Q_i - Q_{min}} \]

です。ブロック内では

• \(0 \le P_i - P_{min} \le K-1\)
• \(0 \le Q_i - Q_{min} \le K - 1\)

が成り立つので、二項係数を用いて展開すれば、 \(R(x)\) は \(O(K^2)\) で求まります（ただし、階乗とその逆数を前計算して二項係数を \(O(1)\) 時間で取得することが必要です）。

その後、 \(R(x) x^{P_{min}} (1+x)^{Q_{min}}\) の展開は FFT を用いてそれぞれ \(O(N \log N)\) で求まります。

これを合計すれば答えが出ます。計算量は \(O(K^2 \times (N/K) + (N/K) \times N \log N)\) ですが、 \(K = \sqrt{N \log N}\) とすれば、 \(O(N \sqrt{N \log N})\) 時間です。

実装例

C++, 981ms, 23 MiB https://atcoder.jp/contests/abc481/submissions/80352937

補足

平方分割系の別解法として、\(Q_i\) の値で平方分割する方法があります。この場合、差分が \(O(1)\) であるという条件は不要で、より広い場面で使えます。計算量は \(O(N \sqrt{N \log N})\) です。

37zigenさんのツイートをご覧ください↓

https://x.com/37zigen/status/2114050869939597675