EDPC A - Frog 1 とほとんど同じ DP で解けますね！

問題へのリンク

問題概要

体力が $N$ であるモンスターが 1 体います。高橋君はモンスターに対し、モンスターの体力が 1 以上残っている限り繰り返し攻撃を行います。

高橋君は 1 回の攻撃で、 $\frac{P}{100}$ の確率でモンスターの体力を 2 減らし、 $1 - \frac{P}{100}$ の確率でモンスターの体力を 1 減らします。

モンスターの体力が 0 以下になるまでに行う攻撃回数の期待値を mod 998244353 で出力してください

制約

$1 \le N \le 2 \times 10^{5}$

解法 (1)：Frog DP

この解法でこの問題に挑むためには、複数の前提知識が必要です。

【前提素養】

EDPC A - Frog 1 と同じ DP と聞いてピンと来ない場合 → この Frog 問題を先に解きましょう！
「期待値を mod 998244353 で求めてください」の意味が分からない場合 → mod 998244353 について学びましょう。

後者について、たとえば次の記事などに詳しく書きました。タイトルの 1000000007 のところは、そろそろ 998244353 にした方がいいかもですね (考え方は特に変わらないです)。

qiita.com

期待値 DP を組む

これらの前提知識は大丈夫として、この問題に挑んでいきましょう。難しいのは「期待値についての DP を立てるってどういうこと！？」という部分だと思います。

回数の期待値の DP を組むときには、

dp[i] = (dp[何か] + (回数)) × (確率) + (dp[何か] + (回数)) × (確率) + ...

のような式になることが多いです。今回の問題に適用してみましょう。まず配列 dp を次のように定義します。

dp[i] ← 体力が $i$ のモンスターの体力を 0 以下にするのに必要な攻撃回数の期待値

このモンスターに対して一発攻撃をします。このとき、ダメージが 2 になるか 1 になるか、確率的に場合が分かれます。よって、場合分けして考えます。

最初のダメージが 2 の場合

この場合、モンスターの体力は $i$ から $i - 2$ になります。ただし、モンスターの体力が 1 の場合は、体力が -1 というように負の値になってしまいます。これは扱いづらいので、0 になるとしましょう。まとめると、モンスターの体力は $\max(i - 2, 0)$ になると考えられます。

また、モンスターの体力が $\max(i - 2, 0)$ になった状態からモンスターの体力が 0 になるまでの攻撃回数の期待値は dp[max(i - 2, 0)] となります。

よって、最初の体力が $i$ の状態から、体力 0 以下にするまでの攻撃回数の期待値は dp[max(i - 2, 0)] + 1 となります。「+1」は最初の攻撃の分です。

最初のダメージが 1 の場合

この場合、モンスターの体力は $i$ から $i - 1$ になります。上と同様に考えて、最初の体力が $i$ の状態から、体力 0 以下にするまでの攻撃回数の期待値は dp[i - 1] + 1 となります。

まとめ

最初のダメージが 2 になる確率は $\frac{P}{100}$ 、ダメージが 1 になる確率は $\frac{100 - P}{100}$ ですので、次のように整理できます。

dp[i] = (dp[max(i-2, 0)] + 1) × $\frac{P}{100}$ + (dp[i-1] + 1) × $\frac{100 - P}{100}$

なお、この式の「+1」のところは式変形してより簡単にできます。

dp[i] = dp[max(i-2, 0)] × $\frac{P}{100}$ + dp[i-1] × $\frac{100 - P}{100}$ + 1

コード

計算量は $O(N)$ となります。

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

// ACL の mod 998244353 を使う
#include <atcoder/modint>
using mint = atcoder::modint998244353;
 
int main() {
    // 入力
    int N, P;
    cin >> N >> P;

    // DP
    vector<mint> dp(N + 1, 0);
    for (int i = 1; i <= N; ++i) {
        dp[i] = dp[max(i - 2, 0)] * mint(P) / 100
                + dp[i - 1] * mint(100 - P) / 100 + 1;
    }
    cout << dp[N].val() << endl;
}

解法 (2)：二項係数を使って直接求める

前提素養として、次の主張が納得できる状態が必要です。数学 IA でもお馴染みの話題です。

【前提素養：成功回数の確率】

表が出る確率が $p$ 、裏が出る確率が $1 - p$ であるようなコインを $n$ 回投げます。表が $k$ 回出る確率は

${}_{n}\mathrm{C}_{r} p^{r} (1 - p)^{n - r}$

と求められます。

場合分け

モンスターの体力が $N$ から $0$ になる過程を次のように場合分けして考えます。

CASE 1：クリティカルヒットが $k$ 回、それ以外が $N - 2k$ 回である結果、モンスターの体力がちょうど 0 になる場合 ( $k = 0, 1, \dots$ )
CASE 2：クリティカルヒットが $k$ 回、それ以外が $N - 2k - 1$ 回やった結果、モンスターの体力が 1 になり、その状態からさらにクリティカルヒットが通ってモンスターの体力が -1 になる場合 ( $k = 0, 1, \dots$ )

ここで、ダメージが 2 になる攻撃をクリティカルヒットと呼んでいます。また、クリティカルヒットの確率を $p (= \frac{P}{100})$ と表しています。

CASE 1

この場合、全体の攻撃回数は $N - k$ 回です。そのようになる確率は

${}_{N - k}\mathrm{C}_{k} p^{k} (1 - p)^{N - 2k}$

となります。

CASE 2

この場合、体力が 1 になるまでの攻撃回数は $N - k - 1$ 回です。そのようになる確率は

${}_{N - k - 1}\mathrm{C}_{k} p^{k} (1 - p)^{N - 2k - 1}$

となります。さらに、体力 1 の状態から追加でクリティカルヒットが発動する確率をかけて、

${}_{N - k - 1}\mathrm{C}_{k} p^{k+1} (1 - p)^{N - 2k - 1}$

となります。全体の攻撃回数は $N-k$ です。

コード

以上をまとめて、期待値は次のように求められます。

$\sum_{k} ({}_{N - k}\mathrm{C}_{k} p^{k} (1 - p)^{N - 2k} * (N - k) + {}_{N - k - 1}\mathrm{C}_{k} p^{k+1} (1 - p)^{N - 2k - 1} * (N - k))$

$p^{k}$ のような冪乗値はあらかじめ前処理で求めることで、計算量は $O(N)$ となりますが、ここでは簡単のため、 ACL の pow を使っています。そのため計算量は $O(N \log N)$ です。

なお、二項係数の求め方は次の記事で解説しています。

drken1215.hatenablog.com

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

// 二項係数を求めるライブラリ
constexpr int MOD = 998244353;
using mint = atcoder::static_modint<MOD>;

// 二項係数を求めるための構造体
template<class T> struct BiCoef {
    vector<T> fact_, inv_, finv_;
    constexpr BiCoef() {}
    constexpr BiCoef(int n) noexcept : 
    fact_(n, 1), inv_(n, 1), finv_(n, 1) {
        init(n);
    }
    constexpr void init(int n) noexcept {
        fact_.assign(n, 1), inv_.assign(n, 1), finv_.assign(n, 1);
        for(int i = 2; i < n; i++){
            fact_[i] = fact_[i-1] * i;
            inv_[i] = -inv_[MOD%i] * (MOD/i);
            finv_[i] = finv_[i-1] * inv_[i];
        }
    }
    constexpr T com(int n, int k) const noexcept {
        if (n < k || n < 0 || k < 0) return 0;
        return fact_[n] * finv_[k] * finv_[n-k];
    }
    constexpr T fact(int n) const noexcept {
        if (n < 0) return 0;
        return fact_[n];
    }
    constexpr T inv(int n) const noexcept {
        if (n < 0) return 0;
        return inv_[n];
    }
    constexpr T finv(int n) const noexcept {
        if (n < 0) return 0;
        return finv_[n];
    }
};
BiCoef<mint> bc;
 
int main() {
    // 入力
    int N, P;
    cin >> N >> P;

    // 二項係数の準備
    bc.init(N + 1);

    // クリティカルヒットの確率、そうでない確率
    mint p = mint(P) / 100;
    mint q = mint(100 - P) / 100;

    // 答え
    mint res = 0;
    for (int k = 0; k <= N; ++k) {
        // CASE 1
        if (N-k*2 >= 0) {
            mint prob = bc.com(N-k, k) * p.pow(k) * q.pow(N-k*2);
            res += prob * (N-k);
        }

        // CAE 2
        if (N-k*2-1 >= 0) {
            mint prob = bc.com(N-k-1, k) * p.pow(k+1) * q.pow(N-k*2-1);
            res += prob * (N-k);
        }
    }
    cout << res.val() << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 280 E - Critical Hit (水色, 500 点)