Candies
解法
$ f_i := 1 + x + x^2 + \ldots + x^{a_i}とおきます。
すると、答えは多項式$ F = \prod_i f_iの$ K次の係数となります。
プログラム上で多項式は、配列で表現できます。
dp[i]が多項式の$ i次の係数を表すものと約束すれば良いです。
多項式の掛け算は、一般にはその次数$ dに対して$ O(d^2)ですが、ゼロでない係数がほとんどないような多項式(疎な多項式)の掛け算は高速に行うことができます。そこで$ f_iを疎にすることを考えます。
等比数列の和の公式により、$ f_i = \frac{1 - x^{a_i + 1}}{1 - x}となります。よって、 $ F = \frac{1}{(1-x)^N} \prod_{i} (1 - x^{a_i + 1})
です。
あとは、この多項式の乗算を実行することに対応するdp遷移を実装します。
多項式$ 1はdp[0] = 1, dp[i] = 0 when i != 0で表現できます。これが「初期状態」となります。
多項式$ 1 - x^{a_i + 1}による乗算は、次のようなdp遷移で表せます。
$ j - (a_i + 1) \ge 0の場合
$ dp(i+1, j) = dp(i, j) - dp(i, j - (a_i + 1))
$ j - (a_i + 1) < 0の場合
$ dp(i+1,j) = dp(i,j)
ここで、iは、「i番目までの多項式$ 1 - x^{a_i + 1}による乗算を行った状態」を表し、
jは多項式の何次の係数なのか、を表します。
つまり、$ dp(i,j)は、「i番目までの多項式$ 1 - x^{a_i + 1}による乗算を行った状態での$ x^jの係数」です。
次に、$ \frac{1}{(1-x)^N}による乗算をします。
maspyさんの記事に書かれていますが、多項式$ 1-xによる割り算は、「元の数列の累積和をとる」操作に対応します。
したがって、元の数列の長さ(=多項式の次数+1)がKであればO(K)で実行可能です。
以上から、前半のdp遷移(計算量はN * Kのオーダー)を行ったあと、累積和をとる操作をN回繰り返し(計算量はN * Kのオーダー)、最後に出てきた多項式の$ x^Kの係数が答えになります。
感想
この問題を解くとき、組み合わせ的な考察をしたのは$ Fの導出までで、あとは多項式の計算を考えるだけで正しいアルゴリズムまでたどり着くことができました。かなり驚きです。やはり魅力的な手法ですね。
参考:他の人の記事