yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

E - Digit Sum Divisible / AtCoder Beginner Contest 336

問題

正整数  n の桁和を、  n を10進法で表したときの各桁の和として定義する。

正整数  n n の桁和で割り切れるとき、  n を良い整数と呼ぶ。

正整数  n が与えられるとき、  n 以下の良い整数の個数を求めよ。

入力

最初の1行に、正整数  n が与えられる。

条件

  • 実行時間制限: 10s
  • メモリ制限: 1024MB
  •  1 \leq n \leq {10}^{14}

出力

答えとなる個数を1行で出力すること。

解法

この問題では、これまでのAtCoderの問題と比較すると実行時間制限が10sとかなり緩めなので、まずは計算量についてあまり考えずに解くことを考えます。

ただし、  n \leq {10}^{14} であるので、  n 以下のすべての整数に対して桁和を計算し、それが良い整数となるのかを判断することはできません。 そこで、  {10}^{14} 以下の整数の桁数は最大で  15 桁であることから、DPを用いて個数を管理することを考えます。

ここで、まずは最終的に得られる整数の桁和の値を固定して考えます。(以下、この桁和の値を  s とします。)

その上で、  \mathrm{dp}_{i, \, j, \, k} を整数について上から  i 桁目までを決定したときに、桁和が  j であり、その整数の  \mathrm{mod} \; s での値が  k である整数の個数として定めます。 (このときに、整数が  n 以下かどうかについては今はまだ考えないこととします。)

現在考えている値の全体の桁数が  d であるとして、このときに上から  (i - 1) 桁目までの  \mathrm{dp} の値が全て決定されているとします。 ここに、 次の  i 桁目として  l が追加されるとします。

このとき、 i - 1 桁目の時点で桁和が  j であった整数は、  i 桁目の時点で桁和が  j + l となります。

また、本来の整数としては元の考えていた整数に対して  l \cdot {10}^{d - i} が加えられることになるので、  i - 1 桁目の時点で  \mathrm{mod} \; s での値が  k である整数は、  i 桁目の時点で  \mathrm{mod}\; s での値が  (k + l \cdot {10}^{d - i}) \, \% \, s となります。

以上から、 \begin{align} \mathrm{dp}_{i, \, j + l, \, (k + l \cdot {10}^{d - i}) \, \% \, s} \leftarrow \mathrm{dp}_{i, \, j + l, \, (k + l \cdot {10}^{d - i}) \, \% \, s} + \mathrm{dp}_{i - 1, \, j, \, k} \end{align} という更新が行えると言えます。

ここで、  i 桁目の値として考えられる値は  0 から  9 までのすべての整数であることから、  0 \leq l \leq 9 です。 従って、すべての  j, \, k に対して  0 \leq l \leq 9 として値を更新することで、  \mathrm{dp}_{i, \, j, \, k} の値が得られるということが言えます。

この問題では、  n \leq {10}^{14} であることから、桁和の最大値として考えられるのは  9 \times 14 = 126 です。 そのため、  i の最大値は  15 であり、  j, \, k の最大値は  126 であることから、この更新を行うことによるおおよその計算回数は、 \begin{align} 15 \times 126 \times 126 \times 10 = 2381400 \end{align} ということになります。

これは  s の値を固定したときの回数であり、  s \leq 126 でもあることから、全体でのおおよその計算回数は, \begin{align} 2381400 \times 126 = 300056400 \fallingdotseq 3 \times {10}^{8} \end{align} ということになります。

この問題の実行時間制限は10sと緩めに設定されていることから、この方法で実行時間制限には間に合いそうということが言えます。

さて、ここからは  n 以下の整数に限る場合、ということについて考えます。

先ほどの  \mathrm{dp} では、考えている値について  n 以下であるという保証がなされていません。 そこで、添字については同じものとして、  \mathrm{dp1} を「考えている値が  n 未満に必ずなるときの個数」とし、  \mathrm{dp2} を「考えている値が  n になりえるときの個数」として定義します。

 n d 桁の整数であるとして、このときに上から  (i - 1) 桁目までの  \mathrm{dp1}, \, \mathrm{dp2} の値が全て決定されているとします。 そして、  n の上から  i 桁目の値が  {a}_{i} であるとします。

このとき、まずは  \mathrm{dp1}_{i - 1, \, j, \, k} についてはその後にどのような整数が  i 桁目として追加されても確実に  n 未満の値となることから、 0 \leq l \leq 9 に対して、先程の  \mathrm{dp} と同様にして、 \begin{align} \mathrm{dp 1}_{i, \, j + l, \, (k + l \cdot {10}^{d - i}) \, \% \, s} \leftarrow \mathrm{dp 1}_{i, \, j + l, \, (k + l \cdot {10}^{d - i}) \, \% \, s} + \mathrm{dp 1}_{i - 1, \, j, \, k} \end{align} という更新が行えます。

また、 i - 1 桁目までは  n になるかもしれない整数についても、  i 桁目として  {a}_{i} 未満の数字が入る場合は確実に  n 未満の整数となります。 従って、  \mathrm{dp2}_{i - 1, \, j, \, k} の値を用いて、  0 \leq l \leq {a}_{i} - 1 を満たすすべての  l に対して、 \begin{align} \mathrm{dp 1}_{i, \, j + l, \, (k + l \cdot {10}^{d - i}) \, \% \, s} \leftarrow \mathrm{dp 1}_{i, \, j + l, \, (k + l \cdot {10}^{d - i}) \, \% \, s} + \mathrm{dp 2}_{i - 1, \, j, \, k} \end{align} という更新が行えます。

最後に、  i - 1 桁目までは  n になるかもしれない整数に対して、  i 桁目として  {a}_{i} が入る場合については、  \mathrm{dp2} の値の更新になるので、 \begin{align} \mathrm{dp 2}_{i, \, j + {a}_{i}, \, (k + {a}_{i} \cdot {10}^{d - i}) \, \% \, s} \leftarrow \mathrm{dp 2}_{i, \, j + {a}_{i}, \, (k + {a}_{i} \cdot {10}^{d - i}) \, \% \, s} + \mathrm{dp 2}_{i - 1, \, j, \, k} \end{align} という更新が行えます。

以上の更新を考えられる  i, \, j, \, k に対して行うことで、  \mathrm{dp1}, \, \mathrm{dp2} が全て得られます。

求めるものは、 n 以下の整数のうち、桁和が  s であり、  s で割り切れるものの総数なので、 n d 桁の整数とすると、  \mathrm{dp1}_{d, \, s, \, 0} + \mathrm{dp2}_{d, \, s, \, 0} となります。 これを考えられるすべての  s に対して導出し、それらの総和が求める答えとなります。

また先程も書きましたが、この問題の実行時間制限は10sと緩めに設定されているので、この方法で実行時間制限に間に合うと言えます。

ソースコード

引数として sum を与えたときに、桁和が sum かつ sum で割り切れる  n 以下の整数の個数を返す関数 solve() を以下のように実装しました。

ll n, dp1[20][150][150], dp2[20][150][150];

ll solve(int sum) {
  // rest は現在見ている桁での値を導出するために用いる
  // digit は上から見たときに現在 10^i の桁を考えているのか、という整数を表す
  // digit_cnt は n の桁数を表す
  ll rest = n, digit = 1, digit_cnt = 1;
  while (digit * 10 <= n) {
    digit *= 10;
    digit_cnt++;
  }

  // 上から0桁目の段階ではまだ n になりえるので、 dp2[0][0][0] = 1 としておく
  dp2[0][0][0] = 1;
  for (ll i = 1; i <= digit_cnt; i++) {
    // dp1 を用いた dp1 の更新
    for (ll j = 0; j < 150; j++) {
      for (ll k = 0; k < 150; k++) {
        for (ll l = 0; l <= 9; l++) {
          dp1[i][j + l][(k + l * digit) % sum] += dp1[i - 1][j][k];
        }
      }
    }

    int l_max = rest / digit; // l_max は先ほどの説明での a[i] に該当する
    // dp2 を用いた dp1 の更新
    for (int j = 0; j < 150; j++) {
      for (int k = 0; k < 150; k++) {
        for (int l = 0; l < l_max; l++) {
          dp1[i][j + l][(k + l * digit) % sum] += dp2[i - 1][j][k];
        }
        // dp2 の更新
        dp2[i][j + l_max][(k + l_max * digit) % sum] += dp2[i - 1][j][k];
      }
    }

    rest %= digit; // rest を 10^i で割った余りを求めることにより、次の位に移行しやすくする
    digit /= 10; // 次の位を考えるために、 digit を 10 で割る
  }

  // 値を返す
  return dp1[digit_cnt][sum][0] + dp2[digit_cnt][sum][0];
}

感想

実行時間制限が10sとなっていたので、最初見たときはビビりました。

整数に対してのDP問題は久しぶりでした。 いつか、EDPC の解説でも触れると思います。