Educational DP Contest
Educational DP Contest - AtCoder
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.
問題
J - Sushi
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.
方針
\(dp[i][j][k]\)として、\(i\)は\(a_i = 1\)の個数、\(j\)は\(a_i = 2\)の個数、\(k\)は\(a_i = 3\)の個数とする。漸化式を作ると、$$dp[i][j][k] = (1+dp[i][j][k])\times \frac{n-i-j-k}{n} + (1+dp[i-1][j][k])\times \frac{i}{n} + \\ (1+dp[i+1][j-1][k])\times \frac{j}{n} + (1+dp[i][j+1][k-1])*\frac{k}{n}$$となる。両辺に\(dp[i][j][k]\)が含まれるので変形すると、$$dp[i][j][k] = \frac{n}{i+j+k}\left(1+dp[i-1][j][k]\times \frac{i}{n}+dp[i+1][j-1][k]\times \frac{j}{n}+dp[i][j+1][k-1]\times \frac{k}{n}\right)$$となり、これを元に三重ループを回していく。
解答
#atcoder template
def main():
import sys
input = sys.stdin.readline
#文字列入力の時は上記はerrorとなる。
#ここにコード
#input
n = int(input())
a = list(map(int, input().split()))
n1, n2, n3 = a.count(1), a.count(2), a.count(3)
#output
dp = [[[0.0] *(n+1) for _ in range(n+1)] for _ in range(n+1)]
#dp[i][j][k] = (1+dp[i][j][k])*(n-i-j-k)/n + (1+dp[i-1][j][k])*i/n
#+(1+dp[i+1][j-1][k])*j/n + (1+dp[i][j+1][k-1])*k/n
#dp[i][j][k] = n/(i+j+k)(1+dp[i-1][j][k]*i/n + dp[i+1][j-1][k]*j/n + dp[i][j+1][k-1]*k/n)
for k in range(n3+1):
for j in range(n2+n3-k+1):
for i in range(n1+n2+n3-j-k+1):
if i == j == k == 0:
continue
temp = n
if i > 0:
temp += i*dp[i-1][j][k]
if j > 0:
temp += j*dp[i+1][j-1][k]
if k > 0:
temp += k *dp[i][j+1][k-1]
dp[i][j][k] = temp/(i+j+k)
print(dp[n1][n2][n3])
#N = 1のときなどcorner caseを確認!
if __name__ == "__main__":
main()
提出結果
Submission #32463663 - Educational DP Contest
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.
感想
dpの初期化のところで、
dp = [[[0] *(n+1) for _ in range(n+1)] for _ in range(n+1)]
ではなく
dp = [[[0.0] *(n+1) for _ in range(n+1)] for _ in range(n+1)]
とすることで早くなる(Pythonの型推論がスムーズになるため)。どっちにしろ\(O(10^6)\)なのでPython3では通らない。
関連リンク
AtCoder
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.
コメント