AOJ 2294 Entangled with Lottery

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2294

解法

ゴールから見るDPでやります．

dp[i][j][k] := 高さ (H から) i まで見た時に，位置 j にいて，それまでに猫が k 本引いているような確率

とします．
高さ 1 から i までに，うさぎがすでに引いている本数を m とすると，猫が引ける高さは残り i - m 箇所です．
今見ている高さがすでにうさぎによって線が引かれていれば，単純に swap するだけです．
引かれていないなら，今見ている高さに猫が線を引く確率 q は，
$\displaystyle q = \frac{_{i - m - 1}C_{K - k - 1}}{_{i - m}C_{K - k}} = \frac{K - k}{i - m}$
となります．
引いた上で，今見ている位置 j に関係する部分に線が引かれる確率 p は，端ならば $\displaystyle p = \frac{1}{N - 1}$，端でなければ $\displaystyle \frac{2}{N - 1}$ です．
あとは，dp[i][j][k] を，「その高さに線をそもそも引かない」「その高さに線を引き，かつそれが j に関係する」「その高さに線を引き，かつそれが j に関係しない」場合の3つに遷移させるだけです．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using pii = pair<int, int>;

int main() {
    cout << fixed << setprecision(10);
    int H, N, P, M, K;
    cin >> H >> N >> P >> M >> K;
    P--;
    vector<int> sw(H, -1);
    for(int i = 0; i < M; ++i) {
        int a, b;
        cin >> a >> b;
        sw[a - 1] = b - 1;
    }

    vector<int> sum(H + 1);
    for(int i = 0; i < H; ++i) {
        sum[i + 1] = sum[i] + (sw[i] != -1);
    }
    vector<vector<double>> dp(N, vector<double>(K + 1));
    dp[P][0] = 1;
    double p = 1.0 / (N - 1);
    for(int i = H - 1; i >= 0; --i) {
        vector<vector<double>> nxt(N, vector<double>(K + 1));
        for(int j = 0; j < N; ++j) {
            for(int k = 0; k <= K; ++k) {
                if(dp[j][k] == 0) {
                    continue;
                }
                if(sw[i] == -1) {
                    double q = (double)(K - k) / (i + 1 - sum[i]); // put possibility
                    if(k == K) {
                        nxt[j][k] += dp[j][k] * (1 - q); // (1 - q) == 1
                    } else {
                        if(j != 0) {
                            nxt[j - 1][k + 1] += dp[j][k] * p * q;
                        }
                        if(j != N - 1) {
                            nxt[j + 1][k + 1] += dp[j][k] * p * q;
                        }
                        nxt[j][k] += dp[j][k] * (1 - q);
                        if(j == 0 || j == N - 1) {
                            nxt[j][k + 1] += dp[j][k] * q * (N - 2) / (N - 1);
                        } else {
                            nxt[j][k + 1] += dp[j][k] * q * (N - 3) / (N - 1);
                        }
                    }
                } else {
                    if(sw[i] == j) {
                        nxt[sw[i] + 1][k] += dp[j][k];
                    } else if(sw[i] + 1 == j) {
                        nxt[sw[i]][k] += dp[j][k];
                    } else {
                        nxt[j][k] += dp[j][k];
                    }
                }
            }
        }
        dp = move(nxt);
    }

    double res = 0;
    for(int i = 0; i < N; ++i) {
        res = max(res, dp[i][K]);
    }
    cout << res << endl;
}