[Python] DP ABC184D

ABC184D

X,Y,Zのいずれかが100のときdp(X,Y,Z)=0 です。そうでないとき、どの硬貨を引いたか3通りの場合を考えることで、

dpの定義：　金貨がX枚、銀貨がY枚、銅貨がZ枚のときの答え

dp漸化式の定義：
$dp(X,Y,Z)$
$=\frac{X}{X+Y+Z}(dp(X+1,Y,Z)+1)+\frac{Y}{X+Y+Z}(dp(X,Y+1,Z)+1)+\frac{Z}{X+Y+Z}(dp(X,Y,Z+1)+1)$

dp終端条件：
$dp(M,M,M)=0$

となります。（金貨を引く確率×金貨を引いた時の操作回数の期待値＋銀貨を……）

この式に従ってDPを行うと答えを求めることができます。実装はメモ化再帰により行うと容易です。

サンプルコード

import sys
sys.setrecursionlimit(10 ** 6)  # 再帰数上限の引き上げ


def dfs(x, y, z):  # メモ化再帰
    if dp[x][y][z] >= 0: # 既に値がわかっている場合はそのまま返す.
        ret = dp[x][y][z]
    else:  # 値がわかっていない場合、漸化式により値を求める.
        ret = 0
        S = x + y + z
        ret += x / S * (dfs(x+1, y, z) + 1)
        ret += y / S * (dfs(x, y+1, z) + 1)
        ret += z / S * (dfs(x, y, z+1) + 1)
        dp[x][y][z] = ret  # メモ
    return ret


X, Y, Z = map(int, input().split())
M = 100
dp = [[[-1] * (M + 1) for _ in range(M + 1)] for _ in range(M + 1)]
# 終端条件（初期化）
for i in range(A, M + 1):
    for j in range(B, M + 1):
        for k in range(C, M + 1):
            if i == M or j == M or k == M:
                dp[i][j][k] = 0

# dp(X, Y, Z)が求める答え
print(dfs(X, Y, Z))