J - Sushi

/icons/hr.icon

出典

EDPC-J

/icons/hr.icon

問題

/icons/hr.icon

考察・解法

期待値 dp といえばと言うほど有名な問題だと思っています。

個人的にまとめておこう思いました。

$ \begin{aligned} {\mathrm dp} [i][j][k] &:= \\ &1 枚の皿が i 枚, 2 枚の皿が j 枚, 3 枚の皿が k 枚残っている時の期待値 \end{aligned}

このように定義します。何番目のお皿がどうこうと言うのが今回は気にする必要がないのがミソです。

遷移を考えます。

したがって、

$ \begin{aligned} {\mathrm dp}[i][j][k] &= dp[i - 1][j][k] \times \dfrac{i}{N} \\ & + {\mathrm dp}[i + 1][j - 1][k] \times \dfrac{j}{N} \\ & + {\mathrm dp}[i][j + 1][k - 1] \times \dfrac{k}{N} \\ & + {\mathrm dp}[i][j][k] \times \dfrac{N - (i + j + k)}{N} \\ & + 1 \end{aligned}

したがって、式変形を行いましょう。

左辺へ持ってくる

$ \begin{aligned} {\mathrm dp}[i][j][k]　- \quad {\mathrm dp}[i][j][k] \times \dfrac{N - (i + j + k)}{N} &= \\ & dp[i - 1][j][k] \times \dfrac{i}{N} \\ & + {\mathrm dp}[i + 1][j - 1][k] \times \dfrac{j}{N} \\ & + {\mathrm dp}[i][j + 1][k - 1] \times \dfrac{k}{N} \\ & + 1 \end{aligned}

共通因数で括る

$ \begin{aligned} {\mathrm dp}[i][j][k] \times \dfrac{i + j + k}{N} &= \\ & dp[i - 1][j][k] \times \dfrac{i}{N} \\ & + {\mathrm dp}[i + 1][j - 1][k] \times \dfrac{j}{N} \\ & + {\mathrm dp}[i][j + 1][k - 1] \times \dfrac{k}{N} \\ & + 1 \end{aligned}

定数を両辺にかける

$ \begin{aligned} {\mathrm dp}[i][j][k] &= \\ & dp[i - 1][j][k] \times \dfrac{i}{i + j + k} \\ & + {\mathrm dp}[i + 1][j - 1][k] \times \dfrac{j}{i + j + k} \\ & + {\mathrm dp}[i][j + 1][k - 1] \times \dfrac{k}{i + j + k} \\ & + \dfrac{N}{i + j + k} \end{aligned}

/icons/hr.icon

コード

code:python

n = int(input())

a = list(map(lambda x: int(x) - 1, input().split()))

# 皿の枚数を計算する

for i in range(n):

for k in range(n + 1):

for j in range(n + 1):

for i in range(n + 1):

p = i + j + k

if p == 0 or p > n:

continue

prob = 0

if i >= 1:

if j >= 1:

if k >= 1:

prob += n

x, y, z = cnt