[Python] 動的計画法 TDPC D

#TDPC D

高校数学の確率問題（場合の数）として解くのは、6=2*3の処理が難しいので、情報数学の手法が適している。
部分構造最適性と部分問題重複性を有するので、動的計画法を検討する。

問題内のデータをそのまま実装するのは難しいので、単純化する必要がある。

積について考えるので、サイコロの目に存在する素因数 2,3,5 の指数のみを考えれば良い。

$10^{18}<2^{60}$ なので、指数として考える値はそれぞれ 60 以下でよいです。

動的計画法の設計は次の通り。

$dp[n][c_2][c_3][c_5]$の定義：
$n$ 回振ったときに、出た目の積に含まれる素因数 2,3,5 の個数がそれぞれ$c_2,c_3,c_5$となる確率

dpの漸化式：

dp[n][nc_2][nc_3][nc_5]=\sum_{i=0}^{n-1}\sum_{c_2\ge nc_2, c_3 \ge nc_3, c_5 \ge nc_5}dp[i][c_2][c_3][c_5]

dp初期条件：
$dp[0][0][0][0]=1$（0 回振ったとき、目の積を 1 とするので、素因数にいかなる素数も含まれない）
$dp[0][][][]=0$

求めるもの：

$D=2^a 3^b 5^c$のとき

\sum_{p_2\ge a, p_3 \ge b, p_5 \ge c}dp[N][p_2][p_3][p_5]

実装は、ボトムアップ方式の配るDPで行う。

サンプルコード

n, d = map(int, input().split())

# サイコロの目を素因数分解すると、2, 3, 5しか現れない
count2 = 0
count3 = 0
count5 = 0
while d % 2 == 0:
    d //= 2
    count2 += 1
while d % 3 == 0:
    d //= 3
    count3 += 1
while d % 5 == 0:
    d //= 5
    count5 += 1

# 2, 3, 5以外があったらだめ
if d != 1:
    print(0)
    exit()

# dp初期化
dp = [[[[0] * (count5 + 1) for i in range(count3 + 1)] for j in range(count2 + 1)] for k in range(n + 1)]
dp[0][0][0][0] = 1

# サイコロの目1~6を素因数分解
d2 = [0, 1, 0, 2, 0, 1]
d3 = [0, 0, 1, 0, 0, 1]
d5 = [0, 0, 0, 0, 1, 0]

# dp[i][c2][c3][c5] := サイコロを i 回振って、2^c2 * 3^c3 * 5^c5 になる確率
for i in range(n):
    # c2,c3,c5および6つの目についてi+1のdpを更新
    for c2 in range(count2 + 1):
        for c3 in range(count3 + 1):
            for c5 in range(count5 + 1):
                for j in range(6):
                    nc2 = min(count2, c2 + d2[j])
                    nc3 = min(count3, c3 + d3[j])
                    nc5 = min(count5, c5 + d5[j])
                    dp[i + 1][nc2][nc3][nc5] += dp[i][c2][c3][c5] / 6

print(dp[n][count2][count3][count5])