はまやんはまやんはまやん

hamayanhamayan's blog

Mod i [AtCoder Beginner Contest 207 E]

https://atcoder.jp/contests/abc207/tasks/abc207_e

前提知識

解説

https://atcoder.jp/contests/abc207/submissions/23805901

高度なDP。慣れていないと簡単ではないと思うのだが、500人も解くんですね…

DPっぽさがあるか?

mod109+7で制約もdp[N][N]っぽい感じがあるので、とりあえずDP方針で考えていくというのが順当な所だろう。
よくあるDPとして、「dp[i] := i番目までは確定している」みたいなものがあるので、それをベースとしてみる。
追加すべき情報として何個に切り分けたかというのを追加する。
これは問題の条件の肝となる情報であり、かつ、何個で切り分けたかだけが重要でそれ以降の遷移に
どこで切り分けたかの情報は必要ないということから添え字として使えそうな感じがする。
これらの点から以下のDPで解けるのではと予想がつく。

dp[i][k] := 先頭からi番目の数までで条件に合うようk組作ったときの組み合わせ

これを構築できれば、dp[N][1]~dp[N][N]の総和が答えになる。
ここまでは比較的作れる人も多いと思う。
問題はこれの更新最適化である。

愚直に考えると…

dp[i][k]から遷移させるときはdp[i+j][k+1]に遷移が可能で、かつA[i+1]+A[i+2]+...+A[i+j]がk+1の倍数である必要がある。
このように遷移を考えていく方式を遷移元を起点としているので配るDPと表現したりするが、更新最適化では逆にして考える、つまり貰うDPの方が都合がいい場合が多い。
貰うDPとは遷移先を起点として考える捉え方である。
今回で言うと、dp[i][k]へ遷移できるのは、j<iかつA[j+1]+A[j+2]+...+A[i]がkの倍数であるdp[j][k-1]からであるという風に言い換える。

数式っぽく書くと、
dp[i][k] = sum_{j<iかつA[j+1]+A[j+2]+...+A[i]がkの倍数であるj} dp[j][k-1]
となる。
今回の問題は、このsum部分をO(1)でやる必要があり、実は方法がある。

条件を言い換える

dp[i][k] = sum_{j<iかつA[j+1]+A[j+2]+...+A[i]がkの倍数であるj} dp[j][k-1]

区間というのは扱いにくく、特に条件としてj<iはしょうがないのだが、区間の特徴にiが入ってしまうと区間計算を毎回する必要があり、あまりうれしくない。
ここで累積和としてB[i] = A[1] + A[2] + ... + A[i]を導入し、かつ、modを使うことで、

A[j+1]+A[j+2]+...+A[i]がkの倍数である

というのを

B[i] mod kとB[j] mod kが一致する

という風に言い換える。ことにする。
これで区間についての条件というよりかは点についての条件に言い換えることができる。

dp[i][k] = sum_{j<iかつA[j]=Aiであるj} dp[j][k-1]

ちょっと文字量が減っただけではという感じもするかもしれないが、このように変換することでメモ配列を用意することができる。

メモ配列tot

以下のようなメモっておくための配列を更新しながらDP更新に用いることにする。

tot[k][mo] := これまで計算を終えたdpテーブルについて、sum_{A[j]を(k+1)で割った余りがmoであるj} dp[j][k]

「これまで計算を終えたdpテーブルで」というのが重要で、dp[i][k]の計算時にこのtot配列を使用すれば、i番目の計算時までのdpテーブルについて参照しているので、
j<iというのが自動的に守られることになる。よって、計算は単にdp[i][k] = tot[k - 1][A[i] % k]とすることができる。
よくよく見てみると、ちゃんと計算できていることが見て取れる。

このメモ配列の構築方法は、dpテーブルの1要素の計算を終えた段階でtotの方へ反映していけばいい。
つまり、dp[i][k]が確定したら、それをtot[k][A[i] % (k + 1)]に足し合わせていけばいい。
すると、今後のk+1組目の確定時にそれが利用される。

実装上の注意として実際にはdp[i][k]が確定した直後にtotも更新してしまうと、同じiでの更新時に影響が出てしまう可能性があるので、
i番目を処理するときはdp[i][k]をすべて計算してから、そのあとにdp[i][k]をtotに足し合わせる作業をするようにすること。

ちょっと分からないんですが…

DP高速化(特に今回みたいな高度なメモ化というか前計算)はコードだけ見るとむっちゃ短いけど、何やってるか分からないことが多いと思う。
自分もなるべく細かく書くように心がけたが、理解が難しかった場合は他にも解説が無いか探して読んでみることを勧める。
同じ解法でも言い方の違いで理解できたりするので、多読してみるといい(公式の動画解説を見れば大体理解できそうではあるけれど)

int N; ll A[3010];
mint dp[3010][3010];  // dp[i][k] := i番目まででk組作った
mint tot[3010][3010]; // tot[k][mo] := k組作っていてこれまでの総和を(k+1)で割った余りがmoである
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> A[i];

    dp[0][0] = 1;
    tot[0][0] = 1;
    ll sum = 0;
    rep(i, 0, N) {
        sum += A[i];
        rep(k, 1, N + 1) dp[i + 1][k] = tot[k - 1][sum % k];
        rep(k, 1, N + 1) tot[k][sum % (k + 1)] += dp[i + 1][k];
    }
    
    mint ans = 0;
    rep(k, 1, N + 1) ans += dp[N][k];
    cout << ans << endl;
}