Atcorder ABC132 F-Small Products

リンク

• https://atcoder.jp/contests/abc132/tasks/abc132_f
• 提出 https://atcoder.jp/contests/abc132/submissions/6257219?lang=ja

まとめ

• dpで解ける問題
• 単純なdpだとテーブルが爆発する
• 数が少ない時は良いが、大きい数に対応するためにも工夫が必要

問題

K個の整数をならべて、隣り合う整数同士を書けた結果がNを下回るような整数列の数を数える。

例えば K=2,N=3のとき

{1,1},{2,1},{3,1},{1,2},{1,3}で5個

方針

dp[i][j]で"i+2"が"並べた整数の個数"、"j"が"今注目している整数に対して何個の整数列ができるか+それまでの合計値"として組む。

#include <iostream>
#include <vector>

int main() {
  int n, k;
  std::cin >> n >> k;

  --k;
  int size = n + 1;
  std::vector<int64_t> dp(size * k);

  for (int64_t i = 0; i < k; ++i) {
    dp[i * size + 0] = 0;
    for (int64_t j = 1; j < size; ++j) {
      int64_t &now = dp[i * size + j];
      const int64_t &left = dp[i * size + j - 1];
      if (i == 0) {
        now = n / j + left;
      } else {
        now = dp[(i - 1) * size + n / j] + left;
      }
    }
  }

  std::cout << *(dp.end() - 1) << std::endl;

  printf("     |");
  for (int j = 0; j < size; ++j) {
    printf("%5d", j);
  }
  std::cout << "\n----------------------\n";
  for (int i = 0; i < k; ++i) {
    printf("%5d|", i + 2);
    for (int j = 0; j < size; ++j) {
      printf("%5d", dp[i * size + j]);
    }
    std::cout << std::endl;
  }

  return 0;
}

実行結果

input 3 2 output 5 => OK

$ ./a.out 
3 2
5
     |    0    1    2    3
----------------------
    2|    0    3    4    5

input 10 3 output 147 => OK

10 3
147
     |    0    1    2    3    4    5    6    7    8    9   10
----------------------
    2|    0   10   15   18   20   22   23   24   25   26   27
    3|    0   27   49   67   82   97  107  117  127  137  147

分析

input 10 3の時を分析する。 dpがどのように変化しているかはテーブル参照。

整数を2つ並べた時に注目する。

     |    0    1    2    3    4    5    6    7    8    9   10
----------------------
    2|    0   10   15   18   20   22   23   24   25   26   27

今回の値が0だった場合は、何をかけても0にしかならないため、整数列は1つも存在しない。

今回の値が1だった場合は、1,2,3,4,5,6,7,8,9,10のどれをかけても良いため、10個の候補が存在する。

今回の値が2だった場合は、1,2,3,4,5のどれをかけても良いため、5個の候補が存在する。 1だった場合のところに加算して、合計15個の候補が存在する。

これを10まで繰り返すことで2つの整数に関して何個の候補があるかを調べることができる。

ここに3つ目の整数を追加しよう

     |    0    1    2    3    4    5    6    7    8    9   10
----------------------
    2|    0   10   15   18   20   22   23   24   25   26   27
    3|    0   27   49   67   82   97  107  117  127  137  147

今回の値が0だった場合は、何をかけても0にしかならないため、整数列は1つも存在しない。

今回の値が1だった場合は、[x ｙ 1]という整数列になる。 yの値は1\~10だったらなんでもいい。 yが1だった場合はxが1\~10まで。 yが2だった場合はxが1\~5まで。 …

そのため[x y 1]のときに作れる整数列はyを1\~10まで試した時の合計値になる。 今回は合算しているので列の長さ2、値10の時の27を見ればいいことになる。

今回の値が2だった場合も同様に考える。 [x y 2] と並んでいて、yに1\~5の値が入る。 そのため列の長さ2、値5の時の22になる。 1の時に合算して49だ。

問題点と改善

一番の問題点は、nの値が大きいとメモリ確保でコケることだ。

  int n, k;
  std::cin >> n >> k;

  --k;
  int size = n + 1;
  std::vector<int64_t> dp(size * k);

入力例3を入れると一発でわかる。

$ ./a.out 
314159265 35
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
中止 (コアダンプ)

Kは100までになっているのでN方向のdpをなんとかしたい。

テーブルから増加幅を確認してみる。

// input 10 3
----------------------
    2|   10    5    3    2    2    1    1    1    1    1
    3|   27   22   18   15   15   10   10   10   10   10

列の長さ2の場合、6\~10が常に1になっている。 列の長さ3の場合、6\~10が常に10になっている。 int(10/6) == int(10/10)は同じになるため、目標値に対する商が同じときは合体させることができそうだ。 この関係はint(10/4) == int(10/5)でも同じことがいえる。

圧縮の方針は、

• int(N/a) == int(N/b)となるa,bをまとめ上げる
• そのときの個数を数えておく
• 数え上げのタイミングで"今回の値で作れる整数列”×"今回の値に圧縮されている数"を積み上げていくようにする。

#include <iostream>
#include <vector>

const int64_t MOD = 1000 * 1000 * 1000 + 7;

int main() {
  int64_t n, k;
  std::cin >> n >> k;

  std::vector<int64_t> cand;
  std::vector<int64_t> num;
  cand.push_back(0);
  num.push_back(0);

  for (int64_t i = 1;;) {
    auto div = n / i;
    cand.push_back(div);
    auto next = n / div + 1;
    num.push_back(next - i);
    i = next;
    if (div == 1) {
      break;
    }
  }

  --k;
  auto size = cand.size();
  std::vector<int64_t> dp(size * k);

  dp[0] = 0;
  for (int64_t j = 1; j < size; ++j) {
    int64_t &now = dp[j];
    const int64_t &left = dp[j - 1];
    now = ((cand[j] * num[j]) % MOD + left) % MOD;
  }

  for (int64_t i = 1; i < k; ++i) {
    dp[i * size + 0] = 0;
    for (int64_t j = 1; j < size; ++j) {
      int64_t &now = dp[i * size + j];
      const int64_t &val = dp[i * size - j];
      const int64_t &left = dp[i * size + j - 1];

      now = ((val * num[j]) % MOD + left) % MOD;
    }
  }

  std::cout << *(dp.end() - 1) << std::endl;

  printf("     |");
  for (int j = 0; j < size; ++j) {
    printf("%5d", j);
  }
  std::cout << "\n----------------------\n";
  for (int i = 0; i < k; ++i) {
    printf("%5d|", i + 2);
    for (int j = 0; j < size; ++j) {
      printf("%5d", dp[i * size + j]);
    }
    std::cout << std::endl;
  }

  return 0;
}

というわけでこのプログラム(の最後のテーブル出力部分を消して)を提出してAC。

10 3の場合で確認しておく。 input 10 3 output 147 => ok

10 3
147
     |    0    1    2    3    4    5
----------------------
    2|    0   10   15   18   22   27
    3|    0   27   49   67   97  147

初期化回りは大きく変更している。

ここで同じ商になる値を探りつつ、その時作れる整数列と重ねることができる値の数を求めている。 こんな変な計算しなくても出せる気がしている。

  for (int64_t i = 1;;) {
    auto div = n / i;
    cand.push_back(div);
    auto next = n / div + 1;
    num.push_back(next - i);
    i = next;
    if (div == 1) {
      break;
    }
  }

ここは1列目(つまり整数が2個の時)の初期化になる。 上の計算のついでに行ってしまえばいいという話もある。

  dp[0] = 0;
  for (int64_t j = 1; j < size; ++j) {
    int64_t &now = dp[j];
    const int64_t &left = dp[j - 1];
    now = ((cand[j] * num[j]) % MOD + left) % MOD;
  }

2列目(3個の整数の列)以降はひたすら合算していく。 valで参照しているのがひとつ前の列の、"今の値に対して何個の整数列が作れるか"になる。 それを条件が重なっている数分掛け算してどんどこ足していく。

  for (int64_t i = 1; i < k; ++i) {
    dp[i * size + 0] = 0;
    for (int64_t j = 1; j < size; ++j) {
      int64_t &now = dp[i * size + j];
      const int64_t &val = dp[i * size - j];
      const int64_t &left = dp[i * size + j - 1];

      now = ((val * num[j]) % MOD + left) % MOD;
    }
  }

もし例3を実行するなら最後のテーブル出力のところは消してから実行したほうがいい。(一敗

numの列とdpが2列あれば計算できるのでメモリはもっと節約できる。