Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

【Educational DP-J Sushi】期待値DPは遷移図を描いて式を立てる

一旦500ACをやめて、DP強化に回っています。わからなかった問題をしっかり分析します。

J - Sushi

問題

皿がN個あり、それぞれに1から3つの寿司が乗っている。N面のサイコロを振り、出た目の皿から寿司を一つ食べる。もし寿司がなかったら何もしない。

寿司が全部なくなるためにサイコロを振る回数の期待値を求めよ。

制約: N=300

方針

状態を寿司がk個(k=1,2,3)乗ってるやつが何個あるかで特定出来る。これをx1,x2,x3とする。

f:id:akiradeveloper529:20190924202512j:plain

期待値DPでは遷移図から期待値の計算式を立てることが第一歩。

この問題では、k個の個数を引く確率をpkとして、

  1. f(x1,x2,x3) = p0 * (1+f(x1,x2,x3)) + p1 * (1+f(x1-1,x2,x3)) + p2 * (1+f(x1+1,x2-1,x3)) + p3 * (1+f(x1,x2+1,x3-1))
  2. A0 = p0 * (1+A0) + p1 * (1+A1) + p2 * (1+A2) + p3 * (1+A3)
  3. A0 = p0+p1+p2+p3 + p0 * A0 + p1 * A1 + p2 * A2 + p3 * A3
  4. (1-p0) * A0 = 1 + p0 * A0 + p1 * A1 + p2 * A2 + p3 * A3

と導くことが出来る。どこでも1+となっているのは、1回サイコロを振ってるから。遷移図を描くことで、自己ループを数式に落とし込むことが出来る。

これをメモ化再帰で実装すればよい。

実装

struct Rec {
    N: usize,
    memo: Vec<Vec<Vec<Option<f64>>>>,
}
impl Rec {
    fn f(&mut self, x1: usize, x2: usize, x3: usize) -> f64 {
        if x1+x2+x3 == 0 {
            return 0.0;
        }

        if self.memo[x1][x2][x3].is_some() {
            return self.memo[x1][x2][x3].unwrap();
        }

        let nf = self.N as f64;
        let p1 = x1 as f64 /nf;
        let p2 = x2 as f64 /nf;
        let p3 = x3 as f64 /nf;
        let p0 = 1.0-(p1+p2+p3);

        let mut res = 0.0;
        res += 1.0;
        if x1>=1 {
            res += p1 * self.f(x1-1,x2,x3);
        }
        if x2>=1 {
            res += p2 * self.f(x1+1,x2-1,x3);
        }
        if x3>=1 {
            res += p3 * self.f(x1,x2+1,x3-1);
        }
        res /= (1.0-p0);

        self.memo[x1][x2][x3] = Some(res);
        return res
    }
}
fn solve() {
    input!{
        N:usize,
        a:[usize;N],
    }

    let mut x1=0;
    let mut x2=0;
    let mut x3=0;
    for i in 0..N {
        match a[i] {
            1 => x1+=1,
            2 => x2+=1,
            3 => x3+=1,
            _ => panic!(),
        }
    }

    let mut rec = Rec {
        N: N,
        memo: dvec![None; 301,301,301],
    };
    println!("{}", rec.f(x1,x2,x3));
}