まあそうだろうなあとは思っていたが時間内に式を計算して実装しきる自信がなくてHackやEに逃げた。

problem

空文字列から始めて確率$\frac{p_a}{p_a + p_b}$で末尾にaを確率$\frac{p_b}{p_a + p_b}$で末尾にbを追加することを繰り返す。 必ずしも連続しない部分列でabの形をしたものの数を$x$とする。 $x$が$k$個以上になった時点で終了する。 終了時の$x$の期待値$\frac{P}{Q} \in \mathbb{Q}$を求め$PQ^{-1} \bmod 10^9+7$を答えよ。

solution

DP。$O(k^2)$。

abの数の増え方を考えると、aが$i$個あって部分列abが$j$個あるような列から始めたときの終了時のabの期待値を$\mathrm{dp}(i, j)$とするのが自然に思い付く。 遷移は$\mathrm{dp}(i, j) = \frac{p_a}{p_a + p_b} \cdot \mathrm{dp}(i + 1, j) + \frac{p_b}{p_a + p_b} \cdot \mathrm{dp}(i, j + i)$。 しかしこの制約だけであるとwell-foundedでないのでDPに落ちない。 問題は$2$点。 ひとつは$i = j = 0$のときに左辺が右辺にも出てくることで、これは$\mathrm{dp}(0, 0) = \mathrm{dp}(1, 0)$なので$i \ge 1$についてのみ求めればよい。 もうひとつは加算回aを足し続けるようとするような場合で、aaa...abを足すような大きな遷移で見てやればよい。

大きな遷移を用いて$i + j \ge k$な場合を$\mathrm{dp}(i, j) = \sum_{l \in \omega} ((j + i + l) (\frac{p_a}{p_a + p_b})^l \frac{p_b}{p_a + p_b})$としてbase casesにする。 一般に$r \in [0, 1)$に対し$x = \sum_{i \in \omega} r^i = 1 + rx = \frac{1}{1 - r}$。 同様に$y = \sum_{i \in \omega} ir^i = \sum_{i \in \omega} r^{i + 1} + \sum_{i \in \omega} ir^{i + 1} = rx + ry = \frac{rx}{1 - r} = \frac{r}{(1 - r)^2}$。 よって、$a = \frac{p_a}{p_a + p_b}, \; b = \frac{p_b}{p_a + p_b}$とおいて、$\mathrm{dp}(i, j) = \sum_{l \in \omega} (j + i + l) a^l b = (j + i) b \sum_{l \in \omega} a^l + b \sum_{l \in \omega} l a^l = (j + i)b\frac{1}{1 - a} + b\frac{a}{(1 - a)^2} = \frac{(j + i)(1 - a)b + ab}{(1 - a)^2}$。

ところでWolframAlphaは便利:

implementation

#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; (i) < int(n); ++ (i))
#define REP_R(i, n) for (int i = int(n) - 1; (i) >= 0; -- (i))
#define REP3R(i, m, n) for (int i = int(n) - 1; (i) >= int(m); -- (i))
using ll = long long;
using namespace std;
template <typename X, typename T> auto vectors(X x, T a) { return vector<T>(x, a); }
template <typename X, typename Y, typename Z, typename... Zs> auto vectors(X x, Y y, Z z, Zs... zs) { auto cont = vectors(y, z, zs...); return vector<decltype(cont)>(x, cont); }

ll powmod(ll x, ll y, ll p) { assert (0 <= x and x < p); assert (0 <= y); ll z = 1; for (ll i = 1; i <= y; i <<= 1) { if (y & i) z = z * x % p; x = x * x % p; } return z; }
ll modinv(ll x, ll p) { assert (x % p != 0); return powmod(x, p - 2, p); }
constexpr ll mod = 1e9 + 7;

struct rational { int num, den; };
rational make_rational(ll num, ll den = 1) { return (rational) { int(num % mod), int(den % mod) }; }
rational operator + (rational a, rational b) { return make_rational(a.num *(ll) b.den + b.num *(ll) a.den, a.den *(ll) b.den); }
rational operator - (rational a, rational b) { return make_rational(a.num *(ll) b.den - b.num *(ll) a.den, a.den *(ll) b.den); }
rational operator * (rational a, rational b) { return make_rational(a.num *(ll) b.num, a.den *(ll) b.den); }
rational operator / (rational a, rational b) { return make_rational(a.num *(ll) b.den, a.den *(ll) b.num); }

int main() {
    // input
    int k, pa, pb; scanf("%d%d%d", &k, &pa, &pb);
    // solve
    rational a = make_rational(pa, pa + pb);
    rational b = make_rational(pb, pa + pb);
    rational c = make_rational(1) - a;
    auto dp = vectors(k + 1, k + 1, rational());
    REP3R (i, 1, k + 1) {
        REP_R (j, k + 1) {
            dp[i][j] = i + j >= k ?
                (make_rational(j + i) * c * b + a * b) / (c * c) :
                a * dp[i + 1][j] + b * dp[i][j + i];
        }
    }
    // output
    rational x = dp[1][0];
    int y = x.num * modinv(x.den, mod) % mod;
    printf("%d\n", y);
    return 0;
}