ARC059 - キャンディーとN人の子供
https://atcoder.jp/contests/arc059/tasks/arc059_c
問題概要
N人の子どもがいる。それぞれ、はしゃぎ度という整数値A[i]があるとする。 この時、それぞれの子どもが飴をB[i]個もらっているとき、全体の活発度は以下のように求められる。 $$ \Pi _ {i = 1} ^ {N} A[i] ^ {B[i]} = A[1] ^ {B[1]} \cdot A[2] ^ {B[2]} \cdot ... $$
始めに飴をC個持っているとする。この飴の全通りの配り方をしたときに、対応する{A[i]}の組での活発度の和を$ f(A[1], A[2], \cdots, A[N]) $とする。
Subtask1
{A[i]}が与えられる。この時のfをmod $ 10 ^ 9+7 $で求めよ。
制約: $ 1 \leq N \leq 400, 1 \leq C \leq 400, 1 \leq A[i] \leq 400 $
Perfect Task
{A[i]}が与えられる代わりに、{A[i]}と{B[i]}が与えられる。以下の式をmod$ 10 ^ 9 + 7 $求める。 $$ \sum _ {x _ 1 = A[1]} ^ {B[1]} \sum _ {x _ 1 = A[2]]} ^ {B[2]} \cdots f(x _ 1, x _ 2, \cdots, x _ n) $$
制約： $ 1 \leq N \leq 400, 1 \leq C \leq 400, 1 \leq A[i] \leq B[i] \leq 400 $
考え方
- Subtaskから考える。これはしばらく考えると、$ dp[i][j] = i人目の子どもまで見た時、j個を飴を配ったときの活発度の和 $ とすれば、畳み込みの遷移式となる。今回はmod $ 10 ^ 9 + 7 $ゆえに、NTTやFFTは使えないものの、$ O(NC ^ 2) $愚直に畳み込みを計算すればできる。
- Perfect Taskの場合、もちろん$ x _ i $自体1つ1つをSubtaskのように計算するのはTLE必至である。この畳み込みの性質をうまくまとめて計算する必要がある。
- 式変形を冷静にしてみると、くくりだせるとわかる。よって前計算して$ O(NC\max(A[i], B[i]) + NC ^ 2) $で解けた。
考察
Subtask
まずはSubtaskを考える。これは、以下のものを計算する。 $ A[1] ^ {B[1]} A[2] ^ {B[2]} \cdots A[N] ^ {B[N]} $ ですべての$ B[i] $の組を試す。これは計算量をひとまず度外視して、以下のようなDP遷移を考えられる。 $$ \begin{eqnarray}{ dp[i][j] := i人目の子どもまで見た時、j個を飴を配ったときの活発度の和 \nonumber \ dp[i + 1][j + k] += dp[i][j] * (A[i]) ^ {k} } \end{eqnarray} $$
これは、以下の形式的べき級数(配った飴の数は次数)の$ x ^ {C} $の係数に該当する。 $$ \Pi _ {i = 1} ^ {N}(\sum _ {j = 0} ^ {C}A[i] ^ {j} x ^ {j}) $$
そもそも、もともとのDPも、よく見ると、$ i $次の係数を$ a[i], b[i] $とすると、
$$ a[i + j] += a[i] \times b[j] $$
の形なので、畳み込み(形式的べき級数の積)の係数であるとわかります。
これは畳み込みなので、FFTで$ O(NC ^ 2) $から$ O(N C \log C) $まで落ちますが、落さなくても時間内に間に合いますし落す必要性はありません。
Perfect Task
Subtaskのまま解くと、ここで計算量が爆発します。ここから計算量を落せるパートを探します。
まず、$ x _ 1, x _ 2, \cdots $を決めた後の$ f $の計算は確実に落とせません。そして$ f $の値を求めるにはこの方法しかありません。
というわけで、どうにかして外のΣの畳み込みの部分(全部試している部分)をまとめて計算できるか？という発想になります。
具体例として、$ A[1] $は2つの候補$ b, c $があるが、$ A[2 \cdots N] $は1つしかない場合を考えます。 この時、形式的べき級数の係数は、$ A[1] $だけ $$ (\sum _ {j = 0} ^ {C}(b ^ {j} + c ^ {j}) x ^ {j}) $$ となります。(すべての組み合わせのパターンについての計算をして足すというのは、畳み込みの定義そのもの。)
このように書けると気づけば、すべての形式的べき級数で、複数個の候補がある場合はそれらの和をかけ合わせることで、結果全通りの計算の和が求まる(畳み込み)とわかる。
ここまでくればこの問題は正解も同然。あとはスムーズにDPで畳み込みの遷移ができるように、前計算で各子供にあたる形式的べき級数の$ j $次の係数が$ A[i] ^ {j} + (A[i] + 1) ^ j + (A[i] + 2) ^ j + \cdots (B[i]) ^ j) $を求めて、Subtaskと同じような遷移式で全体の畳み込みを計算すればよい。
累乗を計算順序やメモ配列を使って丁寧に計算すると、計算量は累積和で$ O(N C \max(A[i], B[i])) $、畳み込みのDP遷移で$ O(N C ^ 2) $となり、時間内に間に合う。
畳み込みの中の畳み込み
この問題のPerfect Taskでは、畳み込みで演算される係数の積自体がまた、畳み込みという構造である。
そもそも畳み込み自体はべき級数の積であり、それを定めるために細部の積によってべき級数の積の全体を計算する演算であった。その細部の積を計算するとき、積でもあるので、この演算自体もある意味畳み込みであると言える。
今回の場合、Subtaskの畳み込みがメインであり、そこでは整数係数の畳み込みが行われたが、Perfect Taskの考察の通り、その時の係数は$ A[i] ^ {j} + (A[i] + 1) ^ j + (A[i] + 2) ^ j + \cdots (B[i]) ^ j) $という形である。 もちろんこれらの値は整数の和なので整数であり、積を計算するときは多項式のように畳み込みを実際に行わない。 しかし、係数をそのように分解すれば、メインの形式的べき級数の係数の積でも畳み込み演算が行われ、それが問題の何重ものΣに相当すると綺麗にわかる。
ACソースコード
最後に
形式的べき級数による数え上げ特訓第1問。とても実りのある問題だった。
