📘

ABC248 C - Dice Sumのメモ

2022/09/23に公開

目についたC問題を解いてみるの回です。
今回はいかにも自力で解けなさそうな、比較的最近の緑diffに近い問題です。
https://atcoder.jp/contests/abc248/tasks/abc248_c

考察(不発)

シグマ記号とか見るとウッとなりますが、よく見ると単純なパターンで、要するに数列の全要素の合計値っぽいです。
それで各要素は1以上M以下、合計値はK以下で、KはNM以下なので、全部1の場合から全部Mの場合まであり得ます。
なので全探索しますとN ^ M通りを探索することになります。両方とも最大値の50の場合、50の50乗で、8881784197001252323389053344726562500000000000000000000000000000000000000000000000000通りです。この世の終わりが来てもまだ終わってなさそうな素敵な数字が出ました。

これをどうやって現実的な時間で終わらせるかですが、まずは特定のAiを固定することを考えてみました。しかし、ダメそうです。AiA{i+1}の間に直接的な大小関係がないため、Aiを固定したところでA{i+1}の範囲が狭まりません。自分の手札ではこれで手詰まりです…

視点を変えて、 998244353 で割った余りを求めるというほうに注目してみます。こういうのは何度か見たことあるので、何かで割った余りを求める場合の典型的な例がありそうです。

それで、計算途中で随時余りを求めることでオーバーフローを防げるよということがわかりました。分配法則から、これは私でもわかる話でした。合計値を全部求めて最後に余りを求めるやり方だと、その合計値がオーバーフローしてしまっている可能性があるので、そうならないように随時余りを求める必要があると。
これは答えを求めるために必要な知識ではあるのですが、計算量を削減できるわけではありません。なので、この線から調べても無理そうでした。

この辺りで、自分が解ける問題ではないと確信しましたので解説を見に行きます。

解説を見て

これは動的計画法(Dynamic Programming, DP)を使って解ける典型例のようでした。DPは名前だけはよく知っていましたがちゃんと勉強したことはなく、DPを使って解く問題に取り組んだのも初めてでした。DP…とうとう会えたな!

公式解説の動画はDPという言葉をできるだけ使わずに一から解説してくれたので非常に助かりました。まあちゃんと理解したか怪しいのですが、DPの雰囲気を体験することができました。

また、一部しか読めてないのですが以下も参考にしました。
https://qiita.com/drken/items/a5e6fe22863b7992efdb

基本的な考え方としては以下のような感じです。
まず、Aiを固定したところでA{i+1}の範囲が狭まらないと書きましたが、逆に言えばAiが何であろうとA{i+1}とそれ以降のパターン数は共通になります。共通部分があるということはそこが計算量削減の余地になると…

実際のやり方としては、二次元のDP表を作ります。解説では行が「要素を決めた個数」、列が「その時点の合計値」、対応する値が「その時点のパターン数」となっていました。
合計値自体を変数で管理するのではなく、配列の添字を合計値と見立ててそれごとにパターン数を保持するのですね。添字を使ってそれに対応する値をそれぞれ管理するテクニックは以前も見たので、けっこうよく使う手なのかもしれません。

単純に樹形図を書くと、「要素を決めた個数」が1増えるとパターン数は単純にM倍になります。しかし、「その時点の合計値」ごとに管理するようにすると、「要素を決めた個数」が1増えた時に「その時点の合計値」が重複するパターンが出てきます。その分を、それが何個あるのかという情報として圧縮してしまえるということですね。

それで、DPではある行の値が決まっていることが前提で、それをもとに次の行の値を決めるのが基本のようです。値が決まっている最後の行に注目して、それをもとに次の行の値を決めるという「配る」やり方と、値が決まっていない最初の行に注目して、前の行をもとに値を決めるという「貰う」やり方とがあるようです。今回の解説では貰う方法で実装していました。

何をどう貰うのかですが、要素としてあり得る値は1~Mまでですので、1増えたことでその合計値になる場合から、M増えたことでその合計値になる場合までのMパターンあります。それらMパターンそれぞれについて既に「その時点のパターン数」が記録されているので、それを全て貰ってきて足せばいいということになります。(この辺りはわかったようなわからないようなという感じなのですが…)

とりあえず、ここまでわかれば実装できます。

提出コード

一応解説は見ずに書きましたが、ほぼそのままなので単に記憶していて書けただけかもしれません。

fun main() {
    val (n, m, k) = readLine()!!.split(" ").map { it.toInt() }
    val dp = MutableList(n + 1) { IntArray(k + 1) }
    val mod = 998244353

    dp[0][0] = 1

    for(x in 1..n) {
        for(y in 1..k) {
            var now = 0
            for(z in 1..m) {
                if(y - z >= 0) {
                    now += dp[x - 1][y - z]
                    now %= mod
                }
            }
            dp[x][y] = now
        }
    }
    var ans = 0
    for(i in 0..k) {
        ans += dp[n][i]
        ans %= mod
    }

    println(ans)
}

ポイントとして「要素を決めた個数」が0の行と「その時点の合計値」が0の列も作るということです。どちらも0の場合は1通りなので、固定で1を入れます。ここが計算の起点となります。

DP表の走査のための二重ループと、「貰う」処理のためのループを加えて三重ループになります。「要素を決めた個数」が0の行と「その時点の合計値」が0の列は見なくていいので、1始まりとなります。行数はNではなくN+1、列数も同様にK+1なので、ゼロオリジンですがそれぞれnとkまで見ていいです。
縦がxで横がyなのはわかりにくいかもしれませんが、解説を踏襲してそうなっています。i、jにしようかとも思いましたが、そうすると次はkになってしまい紛らわしいのでxyにしちゃいました。
「貰う」処理の際、元が存在しない(負の添字になる)場合があるのでそれを回避するための条件分岐があります。

あとは、オーバーフローしないように随時余りをとります。この問題の場合は随時余りをとればIntの範囲に収まるということでした。
余りのとり方が間違っていてバグらせてました。余りに弱すぎ

計算量としては三重ループなのでO(NMK)になります。DP表の走査がO(NK)、「貰う」処理がO(M)です。

感想

とりあえずDPを体験することができたので、今後は「これDPで解けるかも?」って考えることができるようになったかもしれません。まあ実装できるかというとかなり怪しいので、実際に実装できるようになるためには練習が必要そうです。
今回は二次元表になりましたが、もっと増えることもあるようなので全然まだほんの触りだけ体験したという感じです。

なんというか、値を添え字と見立ててまとめて扱うことで、計算量がDP表の大きさの範囲内に収まるようにするのってすごく賢い感じがしました。

それにしても、DPが必要になるのってD問題以降というイメージだったのですが、C問題でも必要になる場合があるのですね。(別解もあるので「必要」ではないかもしれないけど、正攻法になるという意味で)
まあ、C問題としてはかなりdiffが高い問題なので例外的なのかもしれないですが…

Discussion