D - increment of coins

[hxrxchang]

https://atcoder.jp/contests/abc184/tasks/abc184_d

問題概要

袋の中に金貨、銀貨、銅貨がそれぞれA枚、B枚、C枚ある。 その中から1枚取り出し、同じ種類のコインを2枚戻すというのを、袋の中の枚数が100になるまで繰り返す。 試行回数の期待値はいくつか？

解き方

N = 100 なので、A, B, Cの組み合わせは 100 ** 3。 dp[A][B][C]: AがA枚、BがB枚、CがC枚のときの試行回数の期待値 で間に合う。

• A, B, C どれかが100のとき、期待値は0になる。
• i枚目の期待値は i + 1枚目の期待値に依存している

ため、99 -> 0 で遷移していけばよい。

A, B, C = map(int, input().split())

dp = [[[0.0] * 101 for _ in range(101)] for _ in range(101)]

for a in range(99, A-1, -1):
    for b in range(99, B-1, -1):
        for c in range(99, C-1, -1):
            total = a + b + c
            if total == 0:
                continue
            if a < 100:
                dp[a][b][c] += (a / total) * (dp[a + 1][b][c] + 1)
            if b < 100:
                dp[a][b][c] += (b / total) * (dp[a][b + 1][c] + 1)
            if c < 100:
                dp[a][b][c] += (c / total) * (dp[a][b][c + 1] + 1)

print(dp[A][B][C])