はまやんはまやんはまやん

hamayanhamayan's blog

New Year and Arbitrary Arrangement [Good Bye 2017 D]

http://codeforces.com/contest/908/problem/D

以下の操作を行う。

  • 最初は文字列""
  • 確率PA/(PA+PB)で末尾に'a'を追加
  • 確率PB/(PA+PB)で末尾に'b'を追加
  • 文字列の部分列(連続でなくて良い)にabがK個以上含まれれば操作を終える

こうして作られる全てのありうる文字列について、その文字列に含まれるabの個数の期待値を求めよ。

前提知識

解法

http://codeforces.com/contest/908/submission/33784056 (コード内のコメントアウトは間違っているので無視すること)

最初はDPをする。
dp[a][k] := これまでの文字列でaの個数がa個でabの組がk組ある文字列が作られる確率
これをやるが、最初の1手が少し面倒。
dp[0][0] -> dp[1][0]の遷移だけループが存在するので、ここだけ別途計算する。
dp[0][0]でaを追加するとdp[1][0]だが、bを追加するとdp[0][0]のままである。
その為dp[1][0] = sum{i=0..INF}(dp[0][0] * A^i * B)となる。
他の遷移は、'a'を末尾に追加すれば「dp[a + 1][k]+=dp[a][k] * A」、'b'を末尾に追加すれば「dp[a][k+a] += dp[a][k] * B」となるのでこれで逐一やっていく。
 
次はこれを利用して答えを作っていく。
K≦a+kのときのdp[a][k]を考えると、後1回bが出ればabがK個以上となるので、bが出るまでの無限級数を取って考える。
しかし、無限級数を考えるのはdp[K][k]の場合(a=K)だけ考えれば良い。
これはdp[a][k](a<K)の時で無限級数を考えるとダブって場面を数えてしまうからである。
そのため、dp[a][k](a<K)の場合は普通にBが出た時の期待値を計算し、dp[K][k]の場合はbが出るまでの無限級数で期待値を計算してやる。
 
後者は具体的にはsum{i=0..INF}{(k + K + i) dp[K][k] * A^i * B}を計算していく。
sumの中身を2つに分けるたり、定数を外に出したりすると、このサイトの公式の形に落ちるので、計算ができる。

int K, PA, PB;
mint dp[2010][2010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> K >> PA >> PB;

    mint A = mint(PA) / (mint(PA) + mint(PB));
    mint B = mint(PB) / (mint(PA) + mint(PB));

    dp[0][0] = 0;
    dp[1][0] = (mint(1) + B / (mint(1) - B)) * A;
    rep(a, 0, K + 1) rep(k, 0, K) {
        // select a
        dp[a + 1][k] += dp[a][k] * A;
        
        // select b
        dp[a][k + a] += dp[a][k] * B;
    }

    mint ans = 0;
    rep(a, 0, K) rep(k, 0, K) if (K <= k + a) {
        mint kk = mint(k) + mint(a);
        ans += dp[a][k] * B * kk;
    }
    rep(k, 0, K) {
        mint kk = mint(k) + mint(K);
        ans += dp[K][k] * B * (kk / (mint(1) - A) + (A / (mint(1) - A) / (mint(1) - A)));
    }
    cout << ans << endl;
}