Official

Ex - Avoid Square Number Editorial by physics0523


\(N\) 要素のうち特定の \(i\) 要素が平方数である (残りの \(N-i\) 要素はどちらでもよい) ような列の個数は、形式的冪級数の表現を用いて以下のように記述できます。

\(\displaystyle \prod_{j=1}^{K} [x^{E_j}] \frac{1}{(1-x^2)^{i}(1-x)^{N-i}}\)

包除原理を用いることによって、求める場合の数は以下のように記述できます。

\(\displaystyle \sum_{i=0}^{N} (-1)^i {N \choose i} \prod_{j=1}^{K} [x^{E_j}] \frac{1}{(1-x^2)^{i}(1-x)^{N-i}}\)

あとはこれを計算することを考えます。

\(\displaystyle \frac{1}{(1-x^2)^{i}(1-x)^{N-i}}\)\(i\) のみに依存するので、この式を \(i=0,1,\dots,N\) にわたって計算できればよいです。
式を保持する次数を \(D\) ( 今回の問題では \(D=10000\) 程度まで保持すればよいです )として、そのまま計算するなら \(O(ND^2)\) ですが、これを \(O(ND)\) 程度まで高速することを考えます。

まず、 \(\displaystyle \frac{1}{(1-x)^{N}} = (1 + x + x^2 + \dots)^N\) は、累積和を \(N\) 回取ることで求めることができます。
その後、 \(i=k\) の式から \(i=k+1\) の式へと更新していくことを考えます。ここから先は複数の方針があります。

方針 \(1\) : \(\displaystyle \frac{(1-x)}{(1-x^2)}\) を掛ける

まず、 \((1-x)\) を掛ける操作は、隣接項の公差を取ることで実現されます (これは \(\displaystyle \frac{1}{(1-x)} = 1 + x + x^2 + \dots\) を掛ける操作を undo することに対応します)。
その後、 \(\displaystyle \frac{1}{(1-x^2)} = 1 + x^2 + x^4 + \dots\) を掛ける操作は、二項ごとの累積和を取ることで実現されます。

実装例(C++)

方針 \(2\) : \(\displaystyle \frac{(1-x)}{(1-x^2)}\) を変形する

\(\displaystyle \frac{(1-x)}{(1-x^2)} = \frac{(1-x)}{(1-x)(1+x)} = \frac{1}{(1+x)}\) なので、 これを掛ければよいです。
\(\displaystyle \frac{1}{(1+x)} = 1 + (-x) + (-x)^2 + \dots = 1 - x + x^2 - \dots\) を掛ける操作は、累積和を取る過程で符号を交互につけることで実現されます。

実装例(C++)

このように、形式的冪級数を利用すると、一見別物のような \(2\) つの方針が実は等価な操作であることもすぐ分かります。

posted:
last update: