[AtCoder][DP][Educational DP Contest][競プロ]J-Sushi

sushi on black surface AtCoder
Photo by cottonbro on Pexels.com

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.

コメント

タイトルとURLをコピーしました