NOSSの雑記

主に競プロでやったことを書きます。

ABC104 D - We Love ABC

問題リンク

D - We Love ABC

問題概要

文字列SのABC数とは以下の条件をすべて満たす整数の組(i,j,k)の個数です。

  • 1 <= i < j < k <= |S|
  • Si = 'A'
  • Sj = 'B'
  • Sk = 'C'

Sに含まれる'?'をA,B,Cのいずれかに置き換えてつくれるすべての文字列についてABC数の総和を求めてください。ただし、和を109+7で割ったあまりを出力すること。

制約

  • 3 <= |S| <= 105
  • Sは'A','B','C','?'のみから構成される。

方針

数え上げなので何かをまとめたり、固定したりするとうまくいくことがあります。今回はABCの 中央 の'B'を固定して考えるとうまくいきます。

'B'を固定すればそのBを使うようなABC数は、(Bより左側の'A'の個数)*(B より右側の'C'の個数) で求められます。これをすべてのBに対して計算し合計すればよさそうです。区間に対する'A'または'C'の個数は事前に累積和を求めておけば高速に求まります。

次に?が含まれている場合を考えます。この場合も同様にして、 ?-B-C や A-?-C の個数を求めます。ただしABCの組をつくるのに関与しない?は任意の文字をとれるので'?'1つにつき3倍します。つまり使われていない?の個数をXとすれば3Xを乗じます。(この計算はXが大きくなる可能性があるので繰り返し二乗法を用います)

Code

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
const ll MOD = ll(1e9 + 7);
const int MAX_N = int(1e5) + 5;
#define REP(i, n) for (int i = 0; i < n; i++)
#define FOR(i, m, n) for (int i = m; i < n; i++)

ll power(ll x, ll n) {
    ll res = 1;
    while (n > 0) {
        if ((n & 1) == 1) {
            res = res * x % MOD;
        }
        x = x * x % MOD;
        n >>= 1;
    }
    return res;
}

ll N, Q, a[MAX_N]{}, c[MAX_N]{}, q[MAX_N];

int main() {
    string s;
    cin >> s;
    N = s.size();
    REP(i, N) {
        if (s[i] == 'A') a[i + 1]++;
        a[i + 1] += a[i];
        if (s[N - i - 1] == 'C') c[i + 1]++;
        c[i + 1] += c[i];
        if (s[i] == '?') q[i + 1]++;
        q[i + 1] += q[i];
    }
    ll ans = 0;
    Q = q[N];
    FOR(i, 1, N-1) {
        if (s[i] == 'B' || s[i] == '?') {
            ll X = Q - (s[i] == '?' ? 1:0);
            ans += (a[i])*(c[N-i])%MOD*power(3,X)%MOD;
            ans += (q[i])*(c[N-i])%MOD*power(3,X-1)%MOD;
            ans += (a[i])*(q[N]-q[i+1])%MOD*power(3,X-1)%MOD;
            ans += (q[i])*(q[N]-q[i+1])%MOD*power(3,X-2)%MOD;
            ans %= MOD;
        }
    }
    cout << ans << endl;
    return 0;
}

コメント

この問題はDPでも解けるらしいのですがまだ理解してません…