ABC_336_E "Digit Sum Divisible"をrubyで解く

これ解けないと1000位以内にも入れないのか(困惑)。
今回の問題はこちら。

で、公式解説はこちら。

公式解説の解説

1. 公式解説コード12行目。まず「桁和の最終目標」を決める。これは1から126までループします。
2. 公式解説13行目。dp配列は4次元で、

• 一次元目は「上から何桁目か」
• 二次元目は「現在の桁和」
• 三次元目は「(現在の値 / 10^上から何桁目か) % 桁和の最終目標」
• 四次元目は「現在の値はN(値の上限)に張り付いているか」
• 値はパターンの個数
  と言う構造にします。

3. 公式解説14行目。dpに初期値を入れます。Nの最大桁より1桁多い0連打(N = 20なら000、N=2024なら00000)の1通りをOKとします。上から0桁目なので1次元目は0、現在桁和0なので二次元目は0、現在値0なので0 % 桁和の最終目標の3次元目は0、現在の値はNに張り付いている(0桁目を1以上にしたらNを超えてしまうので)ので4次元目は1。よってdp[0][0][0][1] = 1
4. 1次元目、2次元目、3次元目、4次元目でそれぞれループ。さらにその中で「新しく入れる数字」を0から9までループ。
5. 「新しく入れる数字」が来ることで桁和がオーバーしてしまったり、Nを超えてしまっていたら早期リターン。
6. 「次の桁まで埋まった」「古い桁和 ＋ 新しく入れる数字」「(古い剰余 * 10 + 新しく入れる数字) % 最終目標桁和」「ここまでの桁がNに張り付いていて、なおかつ新しく数字を入れてもまだNに張り付いているか」のパターン数 += 「前の桁和まで埋まっている」「古い桁和」「古い剰余」「ここまでの桁がNに張り付いているか」のパターン数
7. 公式解説20行目。「最後の桁まで埋まった」「目的桁和まで到達した」「現在の値 % 桁和 = 0」「Nに等しいパターンとN未満パターンの両方」を合計していく。

...というのが考え方です。rubyで書いたものがこちらです。

TLEじゃん。はいそうです。この問題は10secと制限時間は緩めですが、rubyの速度では間に合いません。ここから高速化する必要があります。

高速化

配列の次元数を落とす

rubyの配列は高機能である代わりに速度に何があります。公式解説は4次元配列をぶん回していますが、1次元目と4次元目を削って二次元配列にしてみます。

Next DP(1次元目)

例えば2次元dp配列を生成する時、1行前しか生成に使わないとします。この場合、2行以上前は保存しておく必要がありません。これから生成する行をnew,生成に必要な数字の入った行をoldと置き、「oldを使ってnewを生成」「newの配列をoldに移動」を規定回数繰り返せば済むわけです。
公式解説の1次元目である桁数ですが、これは1ずつしか進みません。古い桁数と、1桁増えた新しい桁数の2つだけあれば十分です。変数が2つに増える事は許容すれば、1次元目を削除できます。
rubyでdpを解く時には頻出のこのテクニック、私は3年前から使っていましたが、Next DPという正式名称は一昨日知りました。

列数の増加を許容(4次元目)

公式解説の4桁目である「ここまでの桁の時点ではNに張り付いているかどうか」は0と1の2択しかありません。分割しても行数と列数のどっちかが2倍になるだけなので、行数が2倍に増える事を許容して配列の次元数を抑えてしまいましょう。「現在の桁和」を2倍にし、「N上限に張り付いている場合のみ+1」した数値を行インデックスにします。行の数字を2で割った商が現在の桁和、余りがN上限に張り付いているかどうかです。

この2点をクリアした新しいdp配列の構造はこうなります。

• 「上から何桁目か」はdpに持たせない。別の変数に入れておく
• 1次元目は「現在の桁和」　* 2 + 「現在の値はN(値の上限)に張り付いているか」
• 2次元目は「(現在の値 / 10^上から何桁目か) % 桁和の最終目標」
• 値はパターンの個数

早期リターン

正直これが1番効いているような気がします。
上に出てきた公式解説の解説の6にありますが、dpの更新では『「前の桁和まで埋まっている」「古い桁和」「古い剰余」「ここまでの桁がNに張り付いているか」のパターン数』を加算します。逆に言うと、この値が0ならここより内側のループの計算は全部飛ばして次に行ってOKです。

割り算のキャッシュ

https://qiita.com/saka1_p/items/6400e5cedea0b3c286c9

加減乗算に比べて、除算は極めて時間がかかります。一度計算した値は変数に入れておき、二回目以降は読み取るだけにします。

ハッシュのデフォルト値(効果不明)

rubyのhashは「求められたkeyに対応するvalueがない場合、ブロックを評価して返す」という機能を持たせる事ができます。
今回の問題ではループの1番外側で決めた「目標桁和」でmodを取る事になるので、ハッシュのデフォルト値を「key % 目標桁和」に設定します。これで最初のアクセス時にブロックが評価されて除算が回り、二回目以降はキャッシュされた答えが直接返ります。
とは言えrubyはハッシュもブロックも重いので、どこまで効果的なのかは未検証です。

商のパターンが少ない場合は引き算してしまう

キャッシュというのは同じ数値の計算を何度も繰り返す時ほど効率が良くなります。逆に言うと1違うだけの除算を全部計算していたら非効率です。
今回は「10刻みでキャッシュして、その答えに0から9までの数字を足す。その後に目標桁和を超えてしまっていないか確認し、超えてしまっていた場合は目標桁和を一回引く」と言う作業をします。
例えば目標桁和が3なのに、1の位に8とかが来ることとかはあり得ないので二回以上引く必要はありません。

ACコード

https://atcoder.jp/contests/abc336/submissions/49358633

1. 提出コード16行目。「桁和の最終目標」は最大でも9*Nの桁数までなので、126までループする必要はありません。もっとも、maxテストケースでは多分14桁なのでどうせ126まで見ることになるとは思います。
2. 桁和の最終目標が決まったので、剰余のキャッシュ用ハッシュを用意します。
3. dpは二次元。ループ内に押し込んだらlinterに怒られたので、ラムダ式にしておいてループ内でcallします。

• 一次元目は「現在の桁和　* 2 + 現在の値はN(値の上限)に張り付いているか」
• 二次元目は「(現在の値 / 10^上から何桁目か) % 桁和の最終目標」

4. dpに初期値を入れます。現在桁和が0で上限に張り付いているので1次元目は0 * 2 + 1 = 1。現在の値は0なので、剰余は0に決まっていて二次元目は0。よってdp[1][0] = 1
5. 桁を上からループ。new_dpを生成してからdpを一次元目と二次元目でループ。
6. 提出コード25行目。ここで古い方のdpの値がわかるので、0の場合はnextして高速化。
7. 提出コード27行目。古い剰余を10倍してもう一度剰余を取る。計算が重いのでキャッシュしておく。
8. dpの一次元目のindexを分解して「現在の桁和」と「現在の値はN(値の上限)に張り付いているか」を取り出しておく。高速化のためbit演算しているのでわかりにくいが、indexを2で割った時の商と剰余。
9. 「新しく入れる数字」を0から9までループ。
10. 「新しく入れる数字」が来ることで桁和がオーバーしてしまったり、Nを超えてしまっていたら早期リターン。
11. 7で出した剰余に「新しく入れる数字」を足し、目標桁和を超えていたら目標桁和を引く。
12. 現在の桁和に新しく入れる数字を足して、新しい桁和を出す。
13. 新しい桁和を2倍して、まだNに張り付いている場合は1を足す。これでnew_dpの一次元目を算出。
14. 「新しい桁和 * 2 + まだNに張り付いているか」「新しい剰余」のパターン数 += 「古い桁和 * 2 + まだNに張り付いているか」「古い剰余」のパターン数
15. 「最後の桁まで埋まった」「目的桁和まで到達した」「現在の値 % 桁和 = 0」「Nに等しいパターンとN未満パターンの両方」を合計していく。

---

これCとDで引っ掛からなくても無理だっただろうなあ。