F - Sugoroku2 Editorial by lucifer1004


Let’s first simplify this problem to the following form:

  • There is a mission, our succeeding probability is \(p_s\), and the cost is \(E_s\), while our failing probability is \(p_f=1-p_s\), and the cost is \(E_f\).
  • We will stop as long as we succeed, otherwise we keep trying.
  • What is the expected value of our total cost?

The expected value of the total cost can be represented as:

\[ E=E_sp_s+(E_f+E_s)p_fp_s+(2E_f+E_s)p_f^2p_s+\cdots \]

It equals to:

\[ E_sp_s+p_fp_s((E_f+E_s)+(2E_f+E_s)p_f+\cdots) \]

We need to calculate two sums:

\[ \sum_s=E_s+E_sp_f+E_sp_f^2+\cdots=E_s\frac{1}{1-p_f}=\frac{E_s}{p_s} \]

and

\[ \sum_f=E_f+2E_fp_f+3E_fp_f^2+\cdots \]

which is a bit harder.

Actually, we have:

\[ \sum_fp_f=E_fp_f+2E_fp_f^2+\cdots \]

So we have:

\[ (1-p_f)\sum_f=E_f(1+p_f+p_f^2+\cdots)=E_f\frac{1}{1-p_f}=\frac{E_f}{p_s} \]

So

\[ \sum_f=\frac{E_f}{p_s^2} \]

Then we return to the original equation and now we have:

\[ E=E_sp_s+p_fp_s(\sum_f+\sum_s)=E_sp_s+p_fp_s(\frac{E_s}{p_s}+\frac{E_f}{p_s^2})=E_sp_s+(1-p_s)(E_s+\frac{E_f}{p_s}) \]

Now we return to the original problem and try to calculate \(E_s\), \(E_f\) and \(p_s\).

Here, if we arrive at the \(N\)-th position, we succeed. If we are sent back to the \(0\)-th position, we fail.

The idea is similar to LC837 - New 21 Game. We store the sum of the probability of the positions which can come to the current position, then

\[ p_i=\frac{\sum_p}{M} \]

Of course, the broken positions will not contribute to \(\sum_p\). Also, we need to add \(p_i\) (if \(i\) is not broken) and remove \(p_{i-M}\) (if \(i-M\) is not broken) when \(i\geq M\).

But here we also need the expected value \(E_i\), which can be represented as:

\[ E_i=\frac{\sum p_j(E_j+1)}{\sum_p}=1+\frac{\sum p_jE_j}{\sum_p} \]

The expression indicates that instead of calculating \(E_i\), we can calculate \(p_iE_i\) instead:

\[ p_iE_i=\frac{\sum_p}{M}+\frac{\sum p_jE_j}{M} \]

In a similar manner, we can calculate \(p_iE_i\) accumulatively. Note that the broken positions’ \(pE\) contributes to \(p_fE_f\) instead of \(\sum_{pE}\).

Since surpassing the \(N\)-th position is also counted as reaching the \(N\)-th position, our enumeration will stop at the \(N-1\)-th position. For each position that can reach the \(N\)-th position, we will calculate its contribution to \(p_N\) and \(p_NE_N\) according to the proportion \(\frac{i+M-N+1}{M}\).

Finally, we can get \(p_fE_f\), \(p_s=p_N\) and \(p_sE_s=p_NE_N\), and with these values, we can calculate the final answer \(E\).

  • Time complexity is \(\mathcal{O}(N)\).
  • Space complexity is \(\mathcal{O}(N)\).

Source code:

use proconio::input;

fn main() {
    input! {
        n: usize,
        m: usize,
        k: usize,
        a: [usize; k],
    }

    let mut broken = vec![false; n + 1];
    for i in a {
        broken[i] = true;
    }

    let mut prob_exp = vec![0.0; n + 1];
    let mut prob = vec![0.0; n + 1];
    prob[0] = 1.0;

    let mut prob_sum: f64 = 1.0;
    let mut exp_sum: f64 = 0.0;

    if m >= n {
        let goal_pct = (m - n + 1) as f64 / m as f64;
        prob[n] += prob[0] * goal_pct;
        prob_exp[n] += (prob_exp[0] + prob[0]) * goal_pct;
    }

    let mut prob_exp_fail = 0.0;

    for i in 1..n {
        prob[i] = prob_sum / m as f64;
        prob_exp[i] = prob[i] + exp_sum / m as f64;
        if !broken[i] {
            prob_sum += prob[i];
            exp_sum += prob_exp[i];
            if i + m >= n {
                let goal_pct = (i + m - n + 1) as f64 / m as f64;
                prob[n] += prob[i] * goal_pct;
                prob_exp[n] += (prob_exp[i] + prob[i]) * goal_pct;
            }
        } else {
            prob_exp_fail += prob_exp[i];
        }
        if i >= m && !broken[i - m] {
            prob_sum -= prob[i - m];
            exp_sum -= prob_exp[i - m];
        }
    }

    if prob[n] < 1e-8 {
        println!("-1");
    } else if (prob[n] - 1.0).abs() < 1e-8 {
        println!("{}", prob_exp[n]);
    } else {
        let exp_fail = prob_exp_fail / (1.0 - prob[n]);
        println!("{}", prob_exp[n] + (1.0 - prob[n]) * (prob_exp[n] + exp_fail) / prob[n]);
    }
}

posted:
last update: