AtCoder Educational DP Contest：J問題 - うしのおちちの備忘録

AtCoderをやっていて、Educational DP ContestのJ問題に詰まったので復習がてら自分なりにまとめて見ます。数学的に厳密なものではないのでご了承ください。

f:id:kuroneko1259:20190731115328p:plain

文字の定義

${(i, j, k)}$ ：１個の寿司が乗っている皿がi個、２個の寿司が乗っている皿がj個、３個の寿司が乗っている皿がk個ある状態
${dp[i][j][k]}$ ：(i, j, k)の状態からすべて食べる場合の試行回数の期待値
$N$ ：皿の数
${sm = i + j + k}$ ：寿司が乗っている皿の数

めんどうなので，ある状態からすべて寿司を食べるまでに必要な試行回数の期待値のことを、ある状態の期待値と表現します。

考え方

寿司を１個食べるための試行回数の期待値を使って帰納的にときます。
状態(i, j, k)を考えると、

${dp[i][j][k] = ((i, j, k) から１個寿司を食べた後の状態の期待値) + ((i, j, k)から１個寿司を食べるための試行回数の期待値)}$

という漸化式が成り立ちます。これを元に、右辺をそれぞれ求めれば動的計画法で解けそうです。

右辺第１項の計算

(i, j, k)からの遷移先には(i-1, j, k), (i+1, j-1, k), (i, j+1, k-1)の３つがあります。食べる皿は等確率で選ぶので、(i, j, k)の遷移先のうち、それぞれの状態である確率は ${\frac{i}{sm}, \frac{j}{sm}, \frac{k}{sm}}$ となります。よって(i, j, k)から１個寿司を食べていける状態の期待値は

${dp[i-1][j][k] \times \frac{i}{sm} + dp[i+1][j-1][k] \times \frac{j}{sm} + dp[i][j+1][k-1] \times \frac{k}{sm}}$

となります。

右辺第２項の計算

${1 \times \frac{sm}{N} + 2 \times (1 - \frac{sm}{N}) \frac{sm}{N} + 3 \times (1 - \frac{sm}{N})^2 \frac{sm}{N} \cdots}$

となります。（１個の寿司を食べるためにn回かかるということは、n-1回はすべて寿司が乗っていない皿を選び続けなければならない）

一般にqが１未満の時、 ${1 \times p + 2 \times q p + 3 \times q^2 p \cdots = \frac{p}{(1 - q)^2}}$ が成り立つので、(i, j, k)から１個寿司を食べるための試行回数の期待値は ${\frac{N}{sm}}$ となります。これが漸化式の右辺第二項になります。

漸化式

以上を踏まえると漸化式は

${dp[i][j][k] = dp[i-1][j][k] \times \frac{i}{sm} + dp[i+1][j-1][k] \times \frac{j}{sm} + dp[i][j+1][k-1] \times \frac{k}{sm} + \frac{N}{sm}}$

となります。あとはこれを使って動的計画法で解けると思います。私はメモ化再帰で書きました。直感的にメモ化再帰がわかりやすいものと、ループ回す方がわかりやすいものがありませんか．．．？この問題はループ回すやり方の直感的なイメージが湧かなかったのでメモ化再帰をつかいました。

#include <iostream>
#include <vector>
#include <cmath>
#include <iomanip>

using namespace std;

typedef long long ll;

#define INF 1000000009LL

double dp[305][305][305];//dp[i][j][k]:１個の寿司がi個，２個の寿司がj個，３個の寿司がk個あるときすべて食べ終わるまでの試行回数の期待値
int N;

double dfs(int i, int j, int k) {
    if (dp[i][j][k] > -1) return dp[i][j][k];

    double res = 1.0 * N / (i + j + k);
    if (i > 0) res += 1.0 * i / (i + j + k) * dfs(i-1, j, k);
    if (j > 0) res += 1.0 * j / (i + j + k) * dfs(i+1, j-1, k);
    if (k > 0) res += 1.0 * k / (i + j + k) * dfs(i, j+1, k-1);

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

int main() {
    cin >> N;
    int one = 0, two = 0, three = 0;
    for (int i = 0; i < N; i++) {
        int a; cin >> a;
        if (a == 1) one++;
        else if (a == 2) two++;
        else three++;
    }

    for (int i = 0; i <= N; i++) {
        for (int j = 0; j <= N; j++) {
            for (int k = 0; k <= N; k++) {
                dp[i][j][k] = -1;
            }
        }
    }

    dp[0][0][0] = 0;
    cout << fixed;
    cout << setprecision(10) << dfs(one, two, three) << endl;

    return 0;
}

まとめ

Educational DP Contestをアルファベット順に解いていますが、J問題で大きくつまずきました。中級者以上の方はこの問題はさくっと解けてしまうんでしょうか．．．地道に典型的な問題や考え方を身につけていきたいです。