TDPC D - サイコロ - BEN2のブログ

TDPC D - サイコロに関連する、計算量を落とす工夫についてのメモを残します。

参考

Typical DP Contest D - サイコロ - すいバカ日誌

例題

コインを 1 枚投げて、表が出たら 1 点、裏が出たら 2 点を獲得する。これを 4 回行うとき、合計得点が 6 点以上となる確率はいくらか。

求める確率は

$\displaystyle {\left( \frac{1}{2} \right)}^{4} + {}_4 \rm{C} _1 {\left( \frac{1}{2} \right)}^{4} + {}_4 \rm{C} _2 {\left( \frac{1}{2} \right)}^{4}$ $\displaystyle = \frac{11}{16} = 0.6875$

となる。
これを動的計画法 (DP) で求めてみる。

実装例 1 - 素朴な実装

DP テーブルは

${\rm{dp}} [ i ] [ j ] :=$ コインを $i$ 回投げて、合計得点が $j$ 点となる確率

とし、テーブルの列は 9 列 $(0 \leq j \leq 8)$ 用意する。

#include <bits/stdc++.h>
using namespace std;

int main() {
  vector<vector<double>> dp(5, vector<double>(9));
  dp[0][0] = 1.0;
  for (int i = 0; i < 4; i++) {
    for (int j = 0; j < 9; j++) {
      if (dp[i][j] > 0) {
        dp[i+1][j+1] += dp[i][j] / 2.0;
        dp[i+1][j+2] += dp[i][j] / 2.0;
      }
    }
  }
  double ans = 0.0;
  for (int j = 6; j <= 8; j++) ans += dp[4][j];
  cout << ans << endl;
}

実装例 2

TDPC D - サイコロで利用する考え方。
DP テーブルの列は 7 列 $(0 \leq j \leq 6)$ しか用意せず、 $j = 6$ の列は 6 点以上となる確率とする。

#include <bits/stdc++.h>
using namespace std;

int main() {
  vector<vector<double>> dp(5, vector<double>(7));
  dp[0][0] = 1.0;
  for (int i = 0; i < 4; i++) {
    for (int j = 0; j < 7; j++) {
      dp[i+1][min(j+1, 6)] += dp[i][j] / 2.0;
      dp[i+1][min(j+2, 6)] += dp[i][j] / 2.0;
    }
  }
  cout << dp[4][6] << endl;
}

検証

実装例 1 では 7 点、8 点の列も用意したが、実装例 2 ではそれらを 6 点の列にまとめた。
ここで実装例 2 について、モヤッとしないように表にしてみる。

下に示すのは、完成した DP テーブル。気になるのは右下の端っこ。 f:id:BEN2suzuka:20191029123512p:plain DP テーブルの列をわざと $j = 6$ までしか用意していないため、下に示す通り、当然行き場のない矢印が発生してしまう。そこで、 $j = 7$ や $j = 8$ へ向かう矢印については、min 関数を用いて行き先を $j = 6$ に変更する。 f:id:BEN2suzuka:20191029123524p:plain 行き先を変更して、 $j = 7$ や $j = 8$ に加算しようとした値を $j = 6$ に加算することで、 $j = 6$ の列は 6 点以上を獲得する確率を表すことになる。 f:id:BEN2suzuka:20191029123537p:plain テーブルのサイズを節約できただけでなく、答えは $(i, j) = (4, 6)$ の値を出力するだけでよくなった。
このような工夫を TDPC D - サイコロに活かしてみる。

TDPC D - サイコロ

問題 : TDPC D - サイコロ

サイコロを $N$ 回投げて、出た目の積が $D$ の倍数となる確率を求める問題。
制約 : $1 \leq N \leq 100$ かつ $1 \leq D \leq {10}^{18}$

サイコロを $N$ 回投げて、出た目の積を $M$ とすると、必ず $M = 2^{x} \cdot 3^{y} \cdot 5^{z}$ と表せる。
$M$ が $D$ の倍数となるには、 $D$ は 2, 3, 5 以外の素因数を含まない、かつ $D = 2^{a} \cdot 3^{b} \cdot 5^{c}$ とするとき、 $x \geq a$ かつ $y \geq b$ かつ $z \geq c$ であることを満たせばよい。

次に、サイコロを 1 回投げて、どのように状態が推移するかを整理する。
出た目の積が $2^{i} \cdot 3^{j} \cdot 5^{k}$ の状態で、サイコロを 1 回投げる。すると

1 の目が出たとき、 $i, j, k$ は変化なし
2 の目が出たとき、 $i$ が 1 増加
3 の目が出たとき、 $j$ が 1 増加
4 の目が出たとき、 $i$ が 2 増加
5 の目が出たとき、 $k$ が 1 増加
6 の目が出たとき、 $i, j$ がそれぞれ 1 増加

以上の推移となる。

次に、DP テーブルのサイズを考える。

${\rm{dp}} [ n ] [ i ] [ j ] [ k ] :=$ サイコロを $n$ 回投げて、出た目の積が $2^{i} \cdot 3^{j} \cdot 5^{k}$ となる確率

とする。
$0 \leq i \leq 2N$ 、 $0 \leq j \leq N$ 、 $0 \leq k \leq N$ であるから、 $(N+1) \times (2N+1) \times (N+1) \times (N+1)$ の 4 次元テーブルが思い浮かぶ。これでも AC できたが、上の例題でみた考え方を利用する。
$D = 2^{a} \cdot 3^{b} \cdot 5^{c}$ のとき、 $i \geq a$ かつ $j \geq b$ かつ $k \geq c$ となる確率を求めればよいので、サイズは $(N+1) \times (a+1) \times (b+1) \times (c+1)$ で十分で、 ${\rm{dp}} [ n ] [ a ] [ b ] [ c ]$ には、 $n$ 回投げて、2 が $a$ 個以上、3 が $b$ 個以上、5 が $c$ 個以上となる確率を記録するとよい。計算量は $O(Nabc)$ となる。
確認しておくと、 $1 \leq D \leq {10}^{18}$ であるから、大きくても $a$ は 60 近く、 $b$ は 40 近く、 $c$ は 20 近くに収まる。

ソースコードは、Typical DP Contest D - サイコロ - すいバカ日誌が大変参考になりました。この方は 4 次元テーブルではなく、

${\rm{dp}} [ i ] [ j ] [ k ] :=$ $n$ 回目の試行を行う前の状態
${\rm{next}} [ i ] [ j ] [ k ] :=$ $n$ 回目の試行を行った後の状態

この 2 つの 3 次元テーブルを用意されています。とても勉強になります。

ちなみに 4 次元テーブルで実装したソースコードは以下の通り。

#include <bits/stdc++.h>
using namespace std;
using vi = vector<int>;
using vd = vector<double>;
using vvd = vector<vd>;
using vvvd = vector<vvd>;
using vvvvd = vector<vvvd>;

int main() {
  long long N, D; cin >> N >> D;
  vi pf(3);
  vi div = { 2, 3, 5 };
  for (int i = 0; i < 3; i++) {
    while (D % div[i] == 0) {
      pf[i]++;
      D /= div[i];
    }
  }
  if (D != 1) { cout << 0 << endl; return 0; }
  vvvvd dp(N+1, vvvd(pf[0]+1, vvd(pf[1]+1, vd(pf[2]+1))));
  dp[0][0][0][0] = 1.0;
  vi di = { 0, 1, 0, 2, 0, 1 };
  vi dj = { 0, 0, 1, 0, 0, 1 };
  vi dk = { 0, 0, 0, 0, 1, 0 };
  for (int n = 0; n < N; n++) {
    for (int i = 0; i <= pf[0]; i++) {
      for (int j = 0; j <= pf[1]; j++) {
        for (int k = 0; k <= pf[2]; k++) {
          for (int d = 0; d < 6; d++) {
            int ni = min(pf[0], i+di[d]);
            int nj = min(pf[1], j+dj[d]);
            int nk = min(pf[2], k+dk[d]);
            dp[n+1][ni][nj][nk] += dp[n][i][j][k] / 6.0;
          }
        }
      }
    }
  }
  cout << fixed << setprecision(10);
  cout << dp[N][pf[0]][pf[1]][pf[2]] << endl;
}