[Python] Educational DP Contest I – Coins

問題

I – Coins

確率DPの問題。

参考資料

確率 DP を極めよう

回答

TLEだった回答

PyPyではAC。

import sys

# input処理を高速化する
input = sys.stdin.readline

def main():
    N = int(input())
    lst_p = list(map(float, input().split()))

    # dp[i][j] := 最初の i 枚のコインを投げたときに、表が j 枚となる確率
    # 次のコインが表: dp[i+1][j+1] += dp[i][j]∗p[i]
    # 次のコインが裏: dp[i+1][j] += dp[i][j]∗(1.0−p[i])
    dp = [[0] * (N+1) for i in range(N+1)]

    # 初期条件
    dp[0][0] = 1.0

    for i in range(N):
        for j in range(N):
            dp[i+1][j+1] += dp[i][j] * lst_p[i]
            dp[i+1][j] += dp[i][j] * (1.0 - lst_p[i])
    
    ans = 0
    for i in range((N+1)//2, N+1):
        ans += dp[N][i]

    print(ans)

main()

ACだった回答

import sys

# input処理を高速化する
input = sys.stdin.readline

def main():
    N = int(input())
    p = list(map(float, input().split()))

    # dp[j] := 表が j 枚となる確率
    # dp[j] = dp[j-1]*p[i] + dp[j]*(1-p[i]) ただしdp[j]が重複していることに注意
    dp = [0.0] * (N+1)
    # 初期条件
    dp[0]= 1.0

    for i in range(N):
        p_head = p[i]
        p_tail = 1 - p[i]
        dp = [dp[j -1] * p_head + dp[j] * p_tail for j in range(N+1)]
    
    ans = sum(dp[(N//2+1):])

    print(ans)

main()