タピオカオイシクナーレ [yukicoder 995]

https://yukicoder.me/problems/no/995

前提知識

• 行列累乗

解説

https://yukicoder.me/submissions/433443

各タピオカ同士は、互いに影響を及ぼさないので、独立に計算ができる。
各タピオカについて美しさの期待値を求め、その総和をとることで、答えを導こう。

期待値は、美味しさ×タピオカが入る確率で計算ができるので、タピオカが入る確率を求めよう。
もともとタピオカがミルクティーに入っている状態のとき、移動の回数が全体で偶数回数になれば、
入っている状態になる。
これをDPで求めることを考えよう。
dp[i][p] := i回パワーを使ったときの移動回数%2=pであるときの確率
普通に計算すると、iがKなので、TLE
だが、今回は1つ前の値にのみ更新が依存するので、行列累乗が行える。
これで確率が求まったら、あとは期待値を求めていけばいい。

int N, M; ll K; int p, q, B[101010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> M >> K >> p >> q;
    rep(i, 0, N) cin >> B[i];

    mint pon = mint(p) / mint(q);
    mint poff = mint(1) - pon;

    Mat m(2, Vec(2, 0));
    rep(i, 0, 2) rep(j, 0, 2) {
        if (i == j) m[i][j] = poff.get();
        else m[i][j] = pon.get();
    }
    m = fastpow(m, K);

    mint on, off;

    {
        Vec v(2);
        v[0] = 1;
        v = mulMatVec(m, v);
        on = v[0];
    }

    {
        Vec v(2);
        v[1] = 1;
        v = mulMatVec(m, v);
        off = v[0];
    }

    mint ans = 0;
    rep(i, 0, N) {
        if (i < M) ans += mint(B[i]) * on;
        else ans += mint(B[i]) * off;
    }
    cout << ans << endl;
}