[学習記録]AtCoder EDP J-Sushi[Go]

問題

N 枚の皿があります。 皿には 1,2,…,N と番号が振られています。 最初、各 i (1≤i≤N) について、皿 i には ai​ (1≤ai​≤3) 個の寿司が置かれています。

すべての寿司が無くなるまで、太郎君は次の操作を繰り返し行います。

• 1,2,…,N の目が等確率で出るサイコロを振り、出目を i とする。 皿 �i に寿司がある場合、皿 i の寿司を 1 個食べる。 皿i に寿司が無い場合、何も行わない。

すべての寿司が無くなるまでの操作回数の期待値を求めてください。

解法

すべての皿の寿司の数を添え字にdpを持とうとすると、データ数が多くなり、データを読むときも大変になる。
文字列をキーとしたmapを作ってもTLEする
1 0 1 > dp[1][0][1]
1 1 0 > dp[1][1][0]
1 2 3 2 > dp[1][2][3][2]
寿司を食べる順番はないため、1 0 1 と1 1 0は同じ状態とみることができる。そうすると、状態を表すために同じ個数の皿の数をつかえる。
1 2 3 > dp[1][1][1]
1 1 1 > dp[3][0]]0]

dp = make([][][]float64, n+1)
for i := 0; i <= n; i++ {
	dp[i] = make([][]float64, n+1)
	for j := 0; j <= n; j++ {
		dp[i][j] = make([]float64, n+1)
	}
}

皿の個数 i, j, k から一つ食べた時、その時の遷移は
１．i -1, j, k （１個乗っているさらから食べた）
２．i+1, j -1, k （２個乗っている皿から食べた）
３．i, j+1, k -1（３個乗っている皿から食べた）
の３つがある。
最初の皿の数を引数として、すべての皿の数がゼロになるまで再帰的に計算していけば、期待値が求められる。
期待値は（遷移先の期待値の和）＋（その状態の時の期待値）になる
計算量を減らすためにメモしておく

// 最後の１つを食べる時の期待値は、1 / 皿の総数 の逆数になる
dp[1][0][0] = float64(n)

// ゼロ個の皿の数を第一引数、そのあとに各々の皿の数をc[i]として渡している
ans := calculate(0, c[1], c[2], c[3])

func calculate(c0, c1, c2, c3 int) float64 {
    // 計算した結果、期待値がゼロになることは無い
	if dp[c1][c2][c3] != 0 {
		return dp[c1][c2][c3]
	}

    // eは計算された期待値
	dp[c1][c2][c3] = e + e1 + e2 + e3
	return dp[c1][c2][c3]
}

あとはそれぞれの状態について、期待値を計算する

１．c1, c2, c3の状態の時の期待値
ゼロの皿が無いときは必ず１に、あるときには再試行があるので、それを加味する。

// 期待値 = 0がない場合：それ以下の期待値＋１
//            ある場合：それ以下の期待値 + (残っている皿の数/全数)
var e float64
if c0 == 0 {
	e = 1
} else {
	e = float64(n) / float64(n-c0)
}

２．遷移先の期待値
遷移先の期待値に、その皿が選ばれる確率を掛ける。
ただし、皿がゼロの時は遷移できないのでゼロになる。

// 期待値 = 各cが１以上の時に計算可能
e1 := 0.0
if c1 > 0 {
	e1 = float64(c1) / float64(n-c0) * calculate(c0+1, c1-1, c2, c3)
}

// 他も同様

コード（AC)

var dp [][][]float64
var c []int
var n int

func run() interface{} {
	n = readInt()

	c := make([]int, 4)
	for i := 0; i < n; i++ {
		idx := readInt()
		c[idx]++
	}

	dp = make([][][]float64, n+1)
	for i := 0; i <= n; i++ {
		dp[i] = make([][]float64, n+1)
		for j := 0; j <= n; j++ {
			dp[i][j] = make([]float64, n+1)
		}
	}
	dp[1][0][0] = float64(n)

	ans := calculate(0, c[1], c[2], c[3])
	return ans
}

func calculate(c0, c1, c2, c3 int) float64 {
	if dp[c1][c2][c3] != 0 {
		return dp[c1][c2][c3]
	}

	// 期待値 = 0がない場合：それ以下の期待値＋１
	//            ある場合：それ以下の期待値 + (残っている皿の数/全数)
	var e float64
	if c0 == 0 {
		e = 1
	} else {
		e = float64(n) / float64(n-c0)
	}

	// それ以外の期待値 = 各cが１以上の時に計算可能
	e1 := 0.0
	if c1 > 0 {
		e1 = float64(c1) / float64(n-c0) * calculate(c0+1, c1-1, c2, c3)
	}

	e2 := 0.0
	if c2 > 0 {
		e2 = float64(c2) / float64(n-c0) * calculate(c0, c1+1, c2-1, c3)
	}

	e3 := 0.0
	if c3 > 0 {
		e3 = float64(c3) / float64(n-c0) * calculate(c0, c1, c2+1, c3-1)
	}

	dp[c1][c2][c3] = e + e1 + e2 + e3
	return dp[c1][c2][c3]
}