Problem p03169 (python)

Fixed submission:
'''
N枚の皿それぞれにa_i個の寿司があり、等確率で皿を選んだ時その皿の寿司を一つ食べる。
その皿に寿司が無い場合は何も行わない。
すべての皿がなくなるまでの操作回数の期待値を求めよ。

terms:
    1<=N<=300
    1<=a_i<=3

input:
3
1 1 1

output:
    5.5

N<=300: O(N^3)が通る

期待値DP。
同じ寿司の量の皿の数にさえ注目すればいいところが鍵。
※PythonもPyPyも遅くて通らない...
再帰をfor文に書き直せばPyPyで通るはず
'''

from collections import defaultdict,deque
import sys, bisect, math, itertools
from functools import lru_cache
# @lru_cache(maxsize=None)
sys.setrecursionlimit(10**8)
input = sys.stdin.readline
INF = float('inf')
mod = 10**9+7
eps = 10**-7
def inp():
    '''
    一つの整数
    '''
    return int(input())
def inpl():
    '''
    一行に複数の整数
    '''
    return list(map(int, input().split()))


n = inp()
a = inpl()
'''
dp[i][j][k] : 
残り1個の皿がi枚、残り2個の皿がj枚、残り3個の皿がk枚の状態から
すべての寿司をなくすのに必要な回数の期待値
'''
one, two, three = 0, 0, 0
for i in a:
    if i == 1:
        one += 1
    elif i == 2:
        two += 1
    else:
        three+=1
MAX = max(one, two, three)
dp = [[[0.0] * (n + 1) for i in range(n + 1)] for j in range(n + 1)]

dp[0][0][0] = 1.0
for k in range(three+1):
    for j in range(two+three+1):
        for i in range(one+two+three+1):
            if i + j + k == 0 or i + j + k > n:
                continue
            res = 0.0
            if i != 0:
                res += dp[i - 1][j][k] * i
            if j != 0:
                res += dp[i + 1][j - 1][k] * j
            if k != 0:
                res += dp[i][j + 1][k - 1] * k
            res += n
            res *= 1.0 / (i + j + k)
            dp[i][j][k] = res
print(dp[one][two][three]-1.0)
# print(dp[one - 1][two - 1][three - 1])
# print(dp)
# @lru_cache(maxsize=None)
# def rec(i, j, k):
#     if i == 0 and j == 0 and k == 0:
#         return 0.0
    
#     res = 0.0
#     if i != 0:
#         res += rec(i - 1, j, k) * i
#     if j != 0:
#         res += rec(i + 1, j - 1, k) * j
#     if k != 0:
#         res += rec(i, j + 1, k - 1) * k
#     res += n
#     res *= 1.0 / (i + j + k)
#     return res

# print(rec(one, two, three))


'''
dp遷移は
次に引く皿の寿司 = 次の状態 × 確率 として
zero = dp[i][j][k] * (N-i-j-k)/N
one =  dp[i-1][j][k] * i/N
two =  dp[i+1][j-1][k] * j/N
three= dp[i][j+1][k-1] * k/N
dp[i][j][k]= zero + one + two + three + 1
(最後の1で追加しているイメージ)
となる。しかし両辺にdp[i][j][k]があると更新できない。(自己ループ)
よってこれをdp[i][j][k]について解くと
dp[i][j][k] = (one + two + three + 1)*N/(i+j+k)

dp[i][j][k] = (dp[i-1][j][k]*i + dp[i+1][j-1][k]*j + dp[i][j+1][k-1]*k + N)/(i+j+k)

ただし、それぞれが0の時に例外処理が必要。
'''

Buggy submission (verdict: Runtime Error):
'''
N枚の皿それぞれにa_i個の寿司があり、等確率で皿を選んだ時その皿の寿司を一つ食べる。
その皿に寿司が無い場合は何も行わない。
すべての皿がなくなるまでの操作回数の期待値を求めよ。

terms:
    1<=N<=300
    1<=a_i<=3

input:
3
1 1 1

output:
    5.5

N<=300: O(N^3)が通る

期待値DP。
同じ寿司の量の皿の数にさえ注目すればいいところが鍵。
※PythonもPyPyも遅くて通らない...
再帰をfor文に書き直せばPyPyで通るはず
'''

from collections import defaultdict,deque
import sys, bisect, math, itertools
from functools import lru_cache
# @lru_cache(maxsize=None)
sys.setrecursionlimit(10**8)
input = sys.stdin.readline
INF = float('inf')
mod = 10**9+7
eps = 10**-7
def inp():
    '''
    一つの整数
    '''
    return int(input())
def inpl():
    '''
    一行に複数の整数
    '''
    return list(map(int, input().split()))


n = inp()
a = inpl()
'''
dp[i][j][k] : 
残り1個の皿がi枚、残り2個の皿がj枚、残り3個の皿がk枚の状態から
すべての寿司をなくすのに必要な回数の期待値
'''
one, two, three = 0, 0, 0
for i in a:
    if i == 1:
        one += 1
    elif i == 2:
        two += 1
    else:
        three+=1
MAX = max(one, two, three)
dp = [[[0.0] * (three + 1) for i in range(two+three + 1)] for j in range(one+two+three + 1)]

dp[0][0][0] = 1.0
for k in range(three+1):
    for j in range(two+three+1):
        for i in range(one+two+three+1):
            if i == 0 and j == 0 and k == 0:
                continue
            res = 0.0
            if i != 0:
                res += dp[i - 1][j][k] * i
            if j != 0:
                res += dp[i + 1][j - 1][k] * j
            if k != 0:
                res += dp[i][j + 1][k - 1] * k
            res += n
            res *= 1.0 / (i + j + k)
            dp[i][j][k] = res
# one,two,three=max(0,one-1),max(0,two-1),max(0,three-1)
print(dp[one-1][two-1][three-1])
# @lru_cache(maxsize=None)
# def rec(i, j, k):
#     if i == 0 and j == 0 and k == 0:
#         return 0.0
    
#     res = 0.0
#     if i != 0:
#         res += rec(i - 1, j, k) * i
#     if j != 0:
#         res += rec(i + 1, j - 1, k) * j
#     if k != 0:
#         res += rec(i, j + 1, k - 1) * k
#     res += n
#     res *= 1.0 / (i + j + k)
#     return res

# print(rec(one, two, three))


'''
dp遷移は
次に引く皿の寿司 = 次の状態 × 確率 として
zero = dp[i][j][k] * (N-i-j-k)/N
one =  dp[i-1][j][k] * i/N
two =  dp[i+1][j-1][k] * j/N
three= dp[i][j+1][k-1] * k/N
dp[i][j][k]= zero + one + two + three + 1
(最後の1で追加しているイメージ)
となる。しかし両辺にdp[i][j][k]があると更新できない。(自己ループ)
よってこれをdp[i][j][k]について解くと
dp[i][j][k] = (one + two + three + 1)*N/(i+j+k)

dp[i][j][k] = (dp[i-1][j][k]*i + dp[i+1][j-1][k]*j + dp[i][j+1][k-1]*k + N)/(i+j+k)

ただし、それぞれが0の時に例外処理が必要。
'''