Educational DP Contest J - Sushi - 輪郭をなぞるだけのブログ

atcoder.jp

問題

$1$ から $N$ の番号が振られている $N$ 枚の皿があり、皿 $i$ には $a_{i}$ 個の寿司が置かれている。
すべての寿司がなくなるまで、次の操作を繰り返す。

$1, 2, ..., N$ の目が等確率で出るさいころを振り、出目を $i$ とする。皿 $i$ に寿司がある場合、皿 $i$ の寿司を $1$ 個食べる。皿 $i$ に寿司がない場合は何も行わない。

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

制約

$1 \leq N \leq 300$
$1 \leq a_{i} \leq 3$

(考察) → 実装

何も分からん。
最初のサンプルの意味が分からず、無限回試行したときの期待値を求めていることを理解するのにも時間がかかった。

DP で解けるらしい。
状態をどう保存するかだが、皿の数が最大で $300$ 個あり、乗っている寿司の個数は $0$ から $3$ 個のいずれかである。
素直に眺めていては、最大で $4^{300}$ 通り考えられる。
これだと白旗を上げざるを得ないが、今回は皿に載っている寿司が $3$ 個までなので、寿司の個数が同じものを一括りにすることで、この問題は解消される。
すなわち、 $1$ 個の寿司が乗っている $i$ 枚の皿と、 $2$ 個の寿司が乗っている $j$ 枚の皿と、 $3$ 個の寿司が乗っている $k$ 枚の皿があるとする。

次に、漸化式をどういった形で立てるかだが、まず $dp(i, j, k)$ を以下のように定義する。

$dp(i, j, k)$ ≔ 寿司の乗っている皿が $(i + j + k)$ 枚残っているときの、寿司を食べきるまでの操作回数の期待値

$(i, j, k)$ 枚の状態から、さいころを $1$ 回振ったときにおこる遷移は下図のようになる。
f:id:OutLine:20200109223858p:plain
ただし、遷移先に $0$ 未満の値がある場合は、そこには遷移しない。(負の枚数は存在しない。)
例えば、さいころを振って、乗っている寿司が $2$ 個である皿を取ることになった (確率： $\frac{j}{N}$ ) 場合、枚数は $(i, j, k)$ から $(i+1,j-1,k)$ と変化する。
他の遷移も同様である。

この遷移図から、次の漸化式が成り立つ。
$dp(i, j, k) = 1 + dp(i, j, k) \times \frac{N - (i + j + k)}{N} + 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}$

右辺の $1$ は、次の状態に遷移するときに $1$ 回さいころを振るが、そのことを表す。
(この理解に相当な時間がかかったし、この解釈で合っているかもわからない。)
右辺の $1$ 以下は、操作回数の期待値を計算している。(追記あり)

ただし、今回は自己ループがあるために、両辺に $dp(i, j, k)$ があるので、式変形してそれをなくす。
(このままだと、 $dp(i, j, k)$ の更新時に、まだ確定していない $dp(i, j, k)$ を使うことになってまずい。)
計算途中は略。
$dp(i, j, k) = \{ N + dp(i-1, j, k) \times i + dp(i+1, j-1, k) \times j + dp(i, j+1, k-1) \times k\} \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \times \frac{1}{i + j + k}$

ただし、 $i+j+k = 0$ のときのみ、ゼロ除算を回避するために別途処理する。

DP は確定している部分を基にして値を更新していくが、どこから確定しているか分かりにくいので、メモ化再帰で求めていく。
確定している部分を基にしてという部分、DP はマルコフ過程っぽい気がしている。

ソースコードも載せておく。
dp テーブルの初期化のやり方がダサい。
3 次元配列のときに、この形式の初期化やったことないけど。

コード

#include <iostream>
#include <vector>
#include <algorithm>
#include <iomanip>
using namespace std;

int n;  // 入力

// dp[i][j][k] := 皿に乗っている寿司の個数が (1,2,3) の皿の個数が (i,j,k) 枚のときに
// 寿司をすべて食べきるまでの操作回数の期待値
vector<vector<vector<double>>> dp(305, vector<vector<double>>(305, vector<double>(305, -1)));

// メモ化再帰
double memo_rec(int i, int j, int k){
    if(dp[i][j][k] != -1)   return dp[i][j][k];
    if(i == 0 && j == 0 && k == 0)   return dp[0][0][0] = 0;
    
    double res = 0;
    res += n;
    if(i > 0)   res += memo_rec(i-1, j, k) * i;
    if(j > 0)   res += memo_rec(i+1, j-1, k) * j;
    if(k > 0)   res += memo_rec(i, j+1, k-1) * k;
    res *= 1.0 / (i + j + k);

    return dp[i][j][k] = res;
}

int main(){
    cin >> n;
    int cnt1 = 0, cnt2 = 0, cnt3 = 0;
    for(int i = 0; i < n; ++i){
        int a;
        cin >> a;
        if(a == 1)  ++cnt1;
        else if(a == 2) ++cnt2;
        else    ++cnt3;
    }
    cout << fixed << setprecision(10);
    cout << memo_rec(cnt1, cnt2, cnt3) << endl;
}

確率漸化式っぽい。
大学入試の勉強でこういう類の問題も解いてた気もするが、完全に忘却の彼方である。

雑記

寿司食べてえ…

追記 (2020/3/12)

$dp(i, j, k) = 1 + dp(i, j, k) \times \frac{N - (i + j + k)}{N} + 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}$
は、
$dp(i, j, k) = (dp(i, j, k) + 1) \times \frac{N - (i + j + k)}{N} + (dp(i-1, j, k) + 1) \times \frac{i}{N} \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ + (dp(i+1 , j-1, k) + 1) \times \frac{j}{N} + (dp(i, j+1, k-1) + 1) \times \frac{k}{N}$
をまとめた式とも取れる。