ARC102-E Stop. Otherwise...

形式べき級数楽しい

問題

atcoder.jp

解法

サンプルを見ると明らかに解には対称性があるので,\(i \leq K + 1\)の時のみ考える.

解を多項式の係数に帰着する.\(x^n\)の係数が,数字を\(n\)個選んだ時の組み合わせの数という風に.

まず,\(a + b = i\)なる\(a, b\)に関しては,片方選べばもう一方は選べないが,その選んだ方は何回選んでもいいので,\(1 + 2x + 2x^2 + \dots\)を掛けることに対応する.このようなペアは,\(i = 2k\)のとき\(k - 1\)個,\(i = 2k + 1\)のとき\(k\)個ある.

\(2a = i\)なる\(a\)に関しては,それは一度選んだらもう選べないので,\(1 + x\)を掛けることに対応する.このような\(a\)は,\(i = 2k\)の時に1個ある.

どのように選んでも問題ない数に関しては,\(1 + x + x^2 + \dots\)を掛けることに対応する.このような数は\(i = 2k\)のとき\(K - 2k + 1\)個,\(i = 2k + 1\)のとき\(K - 2k\)個ある.

以上より,解は
\[
\begin{cases}
(1 + 2x + 2x^2 + \dots)^{k-1} (1+x) (1 + x + x^2 + \dots)^{K-2k + 1} & \text{if } i = 2k \\
(1 + 2x + 2x^2 + \dots)^{k} (1 + x + x^2 + \dots)^{K-2k} & \text{if } i = 2k + 1 \\
\end{cases}
\]
の\(x^N\)の係数となる.

ここで,\(1 + 2x + 2x^2 + \dots = \frac{1+x}{1-x}\),\(1 + x + x^2 + \dots = \frac{1}{1 - x}\)となるから,上の式は\(i\)の偶奇にかかわらず
\[
\frac{(1 + x)^k}{(1 - x)^{K - k}}
\]
と書き直せる.

分子は二項定理,分母は負の二項定理を用いてそれぞれ展開し,\(x^N\)の係数を求めればいい.計算量は\(O(KN)\)

終わり.下の提出はNTTを使っているんですが,これはなんか勘違いしていてNTTを使わないといけないと思っていた.求めるのは1項だけなのでNTTはいりません.

atcoder.jp


形式的べき級数は頭を使わなくても機械的な計算で答えが出るのでうれしい.