ARMERIA

Rubyと競技プログラミングの話 AtCoderやCodeforcesの問題解説記事が多め。

AtCoder Regular Contest 059 E - キャンディーとN人の子供

お題箱より。

E - Children and Candies

部分点解法から満点解法へのステップアップについて重点的に書きます。

解法

部分点解法

部分点解法は公式解説にもしっかり書かれているので軽く。各  i について  A_{i} = B_{i} なので、ある1つの組  (x_{1}, ..., x_{N}) に対してだけ  f(x_{1}, ..., x_{N}) を求めてください、という問題です。

DPを考えましょう。 i-1 人目の子供までにキャンディーを合計  j 個配る方法を列挙し、それら全てについて「 i-1 人目の子供までの嬉しさの積」を全て足したものを  dp\lbrack i \rbrack\lbrack j \rbrack とします。

この状態から  i 人目の子供に  k 個のキャンディーを配るとします。このときそれぞれの場合における「 i 人目の子供までの嬉しさの積」は、それぞれの「 i-1 人目の子供までの嬉しさの積」に  x_{i}^{k} を掛けたものになります。つまり  dp\lbrack i \rbrack\lbrack j \rbrack に係数  x_{i}^{k} を掛けて  dp\lbrack i+1 \rbrack\lbrack j+k \rbrack に遷移すれば良いです。

全ての  (i, j) のペアからあり得る全ての  k について遷移を行うため、計算量は  O(NC^{2}) となります。これが部分点解法です。

満点解法

満点解法において、全ての  (x_{1}, ..., x_{N}) に対して1つ1つ計算をするのは無理です。上手く数式で処理できないか考えましょう。

イメージを掴むため小さい個数で考えてみましょう。 N=2 とします。また  A_{1} + 1 = B_{1}, A_{2} + 1 = B_{2} とします。つまり  x_{1}, x_{2} はそれぞれ2通りの値を動き、全部で4通りの  (x_{1}, x_{2}) = (A_{1}, A_{2}), (A_{1}, B_{2}), (B_{1}, A_{2}), (B_{1}, B_{2}) について  f(x_{1}, x_{2}) の和を求めることになります。

ここでキャンディーの配り方を1通りだけ考えます。子供  1, 2 にそれぞれ  a_{1}, a_{2} 個のキャンディーを配ったとします。このときの嬉しさの積は  x_{1}^{a_{1}}x_{2}^{a_{2}} と計算されます。

4通りの  (x_{1}, x_{2}) について、この嬉しさの積の値を合計してみましょう。

 \displaystyle A_{1}^{a_{1}}A_{2}^{a_{2}} + A_{1}^{a_{1}}B_{2}^{a_{2}} + B_{1}^{a_{1}}A_{2}^{a_{2}} + B_{1}^{a_{1}}B_{2}^{a_{2}}

この式は分配法則の逆、いわゆる「因数分解」を用いて以下の形に変形できます。

 \displaystyle (A_{1}^{a_{1}} + B_{1}^{a_{1}})(A_{2}^{a_{2}} + B_{2}^{a_{2}})

ここが満点解法のキモです。因数分解をすることによって、子供1に関する部分と子供2に関する部分の積として分離することができました。

これは全てのキャンディーの配り方  a_{1}, a_{2} について成り立ちます。そしてより一般的に、 N がいくつで、 A_{i}, B_{i} がいくつ離れているような場合でも同様の関係が成り立ちます。つまり子供が  N 人いて、それぞれの子供に配るキャンディーの数が  a_{1}, ..., a_{N} である場合を考えると、この配り方に対応する嬉しさの積の値を全ての  (x_{1}, ..., x_{N}) に対して合計したものは

 \displaystyle (A_{1}^{a_{1}} + \cdots + B_{1}^{a_{1}}) \times \cdots \times (A_{N}^{a_{N}} + \cdots + B_{N}^{a_{N}})

と計算することができます。これは  (x_{1}, ..., x_{N}) が1通りである時の嬉しさの積の値

 \displaystyle (x_{1}^{a_{1}}) \times \cdots \times (x_{N}^{a_{N}})

とほぼ同じ形で、各  x_{i}^{a_{i}} がそれぞれ和の形に変わっています。

ここまで整理すると、先ほどの部分点解法で用いたDPとほとんど同じことができることが分かります。つまり  dp\lbrack i \rbrack\lbrack j \rbrack から  dp\lbrack i+1 \rbrack\lbrack j+k \rbrack に遷移するときに、部分点解法では係数  x_{i}^{k} を掛けていましたが、ここが

 \displaystyle (A_{i}^{k} + \cdots + B_{i}^{k})

という係数に変わります。

DPの遷移だけで  O(NC^{2}) 掛かってしまうので、この係数の計算は  O(1) でやりたいです。ここで  1 \le A_{i} \le B_{i} \le 400 であり、 0 \le k \le C \le 400 であることから、正整数の  0 乗累積和、 1 乗累積和、...、 400 乗累積和をそれぞれ  400 以下の範囲まで前計算しておくことで、この係数を  O(1) で引いてくることができます。

ACコード

Submission #10033202 - AtCoder Regular Contest 059