けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 300 E - Dice Product 3 (水色, 500 点)

確率や期待値を計算する DP をするときは、自己ループを除去することが多いね。

EDPC J - ボール」などは類題と言える。先にこの問題を解いておくと、今回の問題も解きやすいかもしれない!

qiita.com

問題概要

 1 以上  6 以下の整数が等確率で出るサイコロがある。

あなたは、今、整数 1 を持っている。この整数が  N 未満である間、次の操作を繰り返す。

  • サイコロを振り、出た目を  x とする
  • 持っている整数に  x を掛ける

最終的に持っている整数が  N に一致する確率を mod 998244353 で求めよ。

制約

  •  2 \le N \le 10^{18}

考えたこと

この手の問題では DP がうまくいくことが多い。素朴に、次のような DP を考えてみる。


dp[i] ← 持っている整数が  i である状態から開始して、整数  N に到達できる確率


このとき答えは dp[1] となる。

dp[i] を表すために、サイコロの目  x の値で場合分けする。

サイコロの目  x その目が出る確率 その後の持っている数 その後  N に到達できる確率
 1  \displaystyle \frac{1}{6}  i dp[i]
 2  \displaystyle \frac{1}{6}  2i dp[2i]
 3  \displaystyle \frac{1}{6}  3i dp[3i]
 4  \displaystyle \frac{1}{6}  4i dp[4i]
 5  \displaystyle \frac{1}{6}  5i dp[5i]
 6  \displaystyle \frac{1}{6}  6i dp[6i]

よって、次のような DP が組める。

dp[i] = dp[i]  \times \displaystyle \frac{1}{6} + dp[2i]  \times \displaystyle \frac{1}{6} + dp[3i]  \times \displaystyle \frac{1}{6} + dp[4i]  \times \displaystyle \frac{1}{6} + dp[5i]  \times \displaystyle \frac{1}{6} + dp[6i]  \times \displaystyle \frac{1}{6}

ただし注意したいことは、左辺と右辺の両方に dp[i] が登場していることだ。このような現象を僕は自己ループと呼んでいる。式変形しよう。まず、dp[i] を左辺に寄せる。

dp[i]  \times \displaystyle \frac{5}{6} = dp[2i]  \times \displaystyle \frac{1}{6} + dp[3i]  \times \displaystyle \frac{1}{6} + dp[4i]  \times \displaystyle \frac{1}{6} + dp[5i]  \times \displaystyle \frac{1}{6} + dp[6i]  \times \displaystyle \frac{1}{6}

よって、次のようになる。

dp[i] = dp[2i]  \times \displaystyle \frac{1}{5} + dp[3i]  \times \displaystyle \frac{1}{5} + dp[4i]  \times \displaystyle \frac{1}{5} + dp[5i]  \times \displaystyle \frac{1}{5} + dp[6i]  \times \displaystyle \frac{1}{5}

必要な DP 状態空間の大きさ

こうして DP が組めたが、 N \le 10^{18} という制約が気になってくる。愚直にやると  O(N) の計算量となって厳しいようにも思える。

しかし、そもそも  i N を割り切らないときは dp[i] = 0 となるのだ。よって、 N の約数であるような  i に対してのみ dp[i] を考えればよい。

そして、 N の約数の個数は高々 103680 個だ。次の記事を参照。

algo-method.com

よって十分間に合う。

コード

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

#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;

long long N;  // 入力
map<long long, mint> dp;  // メモ化用の配列

// メモ化再帰
mint rec(long long n) {
    if (n == N) return mint(1);
    if (dp.count(n)) return dp[n];
    
    mint res = 0;
    for (int x = 2; x <= 6; ++x) {
        if (N % (n * x) == 0) {
            res += rec(n * x) / 5;
        }
    }
    return dp[n] = res;
}

int main() {
    cin >> N;
    cout << rec(1).val() << endl;
}