ARC055 B問題：せんべい - 忘れても大丈夫

問題. せんべい

1 から $N$ までの番号が付けられた $N$ 枚のせんべいが一様ランダムに並べられている．このせんべいを先頭から順番に食べるかどうかを選び最大で $K$ 枚を食べる． $i$ 番目のせんべいを食べるかどうか決めるときに， $i$ 番目のせんべいの番号は分からないが， $i$ 番目のせんべいを含む今まで見てきたすべてのせんべいの番号の大小関係は分かる．番号が $N$ のせんべいを食べた確率が最大となる選び方をしたときの確率を答えよ．

制約： $1 \le K \le N \le 1,000$

解法．確率DP

解説を参考にした．

戦略を考えるときに何が判断材料にできるのかを考えると，何番目のせんべいについて考えているのか，あと何枚食べれるのかという 2 つがある．また，今まで見てきたせんべいの大小関係は分かるが，見たせんべいの番号は分からない．しかし，すべてのせんべいを見たあとならば大小関係から番号 $N$ のせんべいを食べたかどうかは分かる．したがって，今までで見たせんべいで番号が最も大きいものを食べたかどうかという情報も戦略に必要かもしれないので，前に述べた 2 つの情報と合わせて 3 つの情報をもとに動的計画法を行う．

漸化式 dp[i][j][b] を，残り $j$ 枚食べれて，1 から $i$ 番目の中での最大番号のせんべいを食べたかどうかをブール値 $b$ で表して， $i$ 番目のせんべいを見ているときに最適な戦略を行ったときの，番号 $N$ のせんべいを食べた最大確率とする．
$i$ 番目のせんべいを見ているとき， $i$ 番目のせんべいが今まで見たせんべいの中で番号が最大かどうかは分かっており，これらの起こりうる事象は互いに排反である．最大のものを食べた確率は $\frac{1}{i}$ で，食べていない確率は余事象の確率で $1 - \frac{1}{i}$ である．これらの排反事象それぞれの場合について $i$ 番目のせんべいを食べるかどうかの 2 通りの戦略を考えてその中で確率が最大となるものを選ぶことにする．

メモ化再帰で実装を行っている．実際の遷移はソースコードを参照． $i$ 番目のせんべいが最大番号でないときは食べる必要はないので省略することができる．

計算時間： $O(N K)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

int n, k;
vector<vector<vector<double>>> p;

double dfs(const int i, const int remain, const bool can_eat) {
    if (n < i) return can_eat ? 1.0 : 0.0;

    double &res = p[i][remain][can_eat];
    if (0 <= res) return res;

    // 今までの中で最大ではない
    res = (i - 1.0) / i * dfs(i + 1, remain, can_eat);
    if (0 < remain)
        res = max(res, (i - 1.0) / i * dfs(i + 1, remain - 1, can_eat));

    // 今までの中で最大
    double pp = 1.0 / i * dfs(i + 1, remain, false);
    if (0 < remain)
        pp = max(pp, 1.0 / i * dfs(i + 1, remain - 1, true));
    res += pp;

    return res;
}

double canEatProbability() {
    cin >> n >> k;

    using d = double;
    p.resize(n + 1, vector<vector<d>>(k + 1, vector<d>(2, -1.0)));

    return dfs(1, k, false);
}

int main() {
    cin.tie(0); ios::sync_with_stdio(false);
    cout << setprecision(8) << setiosflags(ios::fixed);

    cout << canEatProbability() << endl;

    return 0;
}