roiti46's blog

主に競プロ問題の解説を載せてます

yukicoder No.108 トリプルカードコンプ

うーん、似たような問題を解いたことがあるのに思いつけなかった。

問題はこちら

問題

N種類のカードがあり、新しく1枚カードを買うと、N種類のいずれかが等確率で手に入る。今持っているカードの情報を与えられるので、すべての種類について3枚以上そろえようとするときに必要な購入カード数の期待値を答える。

解法

あと1枚、2枚、3枚必要なカードの種類数を抱えてメモ化再帰やDPをすれば解ける。 kmjpさんのブログを参考にしたため、回答がほとんど同じものになってしまったいる。解説に感謝。

def dp(n1,n2,n3):
    if memo[n1][n2][n3] > -1: return memo[n1][n2][n3]
    
    n123 = n1+n2+n3
    memo[n1][n2][n3] = N*1.0/n123
    if n1 > 0: memo[n1][n2][n3] += dp(n1-1,n2  ,n3  )*n1/n123
    if n2 > 0: memo[n1][n2][n3] += dp(n1+1,n2-1,n3  )*n2/n123
    if n3 > 0: memo[n1][n2][n3] += dp(n1  ,n2+1,n3-1)*n3/n123
    return memo[n1][n2][n3]
    
N = int(raw_input())
A = map(int,raw_input().split())
memo = [[[-1]*101 for i in xrange(101)] for j in xrange(101)]
memo[0][0][0] = 0.0
need = [0]*3
for a in A:
    if a == 0: need[2] += 1
    if a == 1: need[1] += 1
    if a == 2: need[0] += 1
print dp(*need)

典型問題の1つだろうと思うので、きっちりと解けるようになりたい。