EDPC M - Candies

考え方

ABCでいうと、E問題級。

一見、子供を $1$ 人ずつ増やしながら、飴の個数ごとの配り方を継承するだけの動的計画法。

しかし、マスの数が最大で $100 \times 10^5 = 10^7$ 個ある。
そして、$1$ つのマスを埋めるのに最大 $10^5$ 個の数を足し算しないといけない。
つまり、愚直にやると TLE する。

ということで、最大 $10^5$ 個の数を足し算しなくて済むように、sliding window 法を使う。
以上……ではあるのだが、$2$ つのアルゴリズムのコードが絡まる都合で、実装がとてもややこしい。

計算量は $O(NK)$。

入力例1での動作

入力を受け取る。

n: 3
k: 4
a: {1, 2, 3}

合計 $0$ 個から $4$ 個までについて、ここまでの子供たちへの配り方の数を考える。

まだ誰にも配っていないときは、合計 $0$ 個だけが $1$ 通りである。

{1, 0, 0, 0, 0}

$1$ 人目には $0$ 個または $1$ 個を渡せる。

{1, 1, 0, 0, 0}

$2$ 人目には $0$ 個以上 $2$ 個以下を渡せる。

例えば合計 $2$ 個にするには、それまでの合計が $2,1,0$ 個の場合が使える。
方法数は $0+1+1=2$ 通りである。

全て求めると、次のようになる。

{1, 2, 2, 1, 0}

$3$ 人目には $0$ 個以上 $3$ 個以下を渡せる。

更新前の、合計 $0$ 個から $4$ 個までの方法数を $1,2,2,1,0$ とする。

合計 $x$ 個にする方法数は、更新前の合計 $x,x-1,x-2,x-3$ 個の方法数の和である。

これを毎回最初から足すと時間がかかるので、重なっている部分を sliding window で使い回す。

例えば合計 $4$ 個について、$3$ 人目に $1$ 個以上 $3$ 個以下を渡す場合の和は $2+2+1=5$ である。
$0$ 個渡す場合の $0$ 通りも加えて、合計 $5$ 通りとなる。

次に合計 $3$ 個を考える。
先ほどの和 $5$ から、合計 $3$ 個の $1$ 通りを外す。
代わりに、合計 $0$ 個の $1$ 通りを加えると $5-1+1=5$ となる。

$0$ 個渡す場合の $1$ 通りも加えるので、合計 $6$ 通りとなる。

同じように窓をずらすと、次のようになる。

更新後は次のようになる。

{1, 3, 5, 6, 5}

飴を合計 $4$ 個配る方法は $5$ 通りである。
したがって、答えは $5$。

注意点

答えは $10^9+7$ で割った余りを要求されているので、剰余類環の考えに従って処理する。
足し算や引き算をするたびに結果を % 1000000007 する。
引き算で負になることがあるので、出力前に負数を補正する。

別解

特になし。