Typical DP Contest - D サイコロ

D: サイコロ - Typical DP Contest | AtCoder

  • 最初素因数分解ライブラリ関数をつかって D を分解しようとして TLE していた。
template <typename T>
using V = vector<T>;
template <typename T>
using VV = V<V<T>>;
using Vd = V<double>;


class d_saikoro {
public:
    void solve(void) {
            int N;
            ll  D;
            cin>>N>>D;

            int e2 = 0;
            int e3 = 0;
            int e5 = 0;

            // D はかなり大きいが冪数はずっと小さくてすむ
            while (D%2 == 0) { ++e2; D /= 2; }
            while (D%3 == 0) { ++e3; D /= 3; }
            while (D%5 == 0) { ++e5; D /= 5; }

            // サイコロに含まれない素数が入っているとダメ
            if (D > 1)
            {
                cout<<0<<endl;
                return;
            }

            // dp[t][i][j][k] := 2がi, 3がj ,5がk個出る確率
            VV<V<Vd>> dp(2, VV<Vd>(e2+1, V<Vd>(e3+1, Vd(e5+1, 0))));
            dp[0][0][0][0] = 1.0;  // 0 回の出現確率は 1.0

            // O(N*e2*e3*e5) <= O(N*e2*e3*(N-(e2+e3))) = 3702600 <= 4*10^6
            // ※ i,j,k を N ループではなく e2,e3,e5 loop させることで N^4 loop を回避している
            int cur = 0;
            REP(t, N)
            {
                int next = cur^1;
                REP(i, e2+1) REP(j, e3+1) REP(k, e5+1) dp[next][i][j][k] = 0.0;

                REP(i, e2+1)
                REP(j, e3+1)
                REP(k, e5+1)
                {
                    dp[next][min(i+1, e2)][j][k] += dp[cur][i][j][k]*1.0/6.0; // 2
                    dp[next][min(i+2, e2)][j][k] += dp[cur][i][j][k]*1.0/6.0; // 4
                    dp[next][i][min(j+1, e3)][k] += dp[cur][i][j][k]*1.0/6.0; // 3
                    dp[next][i][j][min(k+1, e5)] += dp[cur][i][j][k]*1.0/6.0; // 5
                    dp[next][min(i+1, e2)][min(j+1, e3)][k] += dp[cur][i][j][k]*1.0/6.0; // 6
                    // (i,j,k) -> (i,j,k) という遷移もあるのでここも +=
                    dp[next][i][j][k] += dp[cur][i][j][k]*1.0/6.0;            // 1
                }
                cur = next;
            }
            cout<<dp[cur][e2][e3][e5]<<endl;
    }
};