AGC029 A - Irreversible operation

問題

• AGC029 A - Irreversible operation

自分の考察

書き出してみたら最終的にどんな文字列も、

WWW...BBB..

になることが分かった。
よって各 W の後ろに B が何個あるのか数えてその和を取れば良い。

一回一回和を取っていると TLE になると思い、累積和を考えた。 実装は、

1. Wのインデックスを取得
2. 各Wの間に何個のBが存在するかを求める
3. それらの累積和を取る

というもの。しかしWA。なぜWAになるのかが分からずそのままコンテスト終了。WAになったソースコードは以下。

#include <iostream>
#include <cstdio>
#include <string>
#include <vector>
#include <numeric>  // accumulate, partial_sum
using namespace std;
 
#define REP(i, n) for (int i=0; i < (int)(n); i++)  // 0 ~ n-1
#define SUM(vec) accumulate((vec).begin(), (vec).end(), 0)  // 0 は初期値
typedef long long ll;
 
vector<ll> findChar(string& str, char pattern) {
    size_t len = str.length();
    vector<ll> vec;
    REP(i, len) {
        if (str[i] == pattern) {
            vec.push_back(i);
        }
    }
    return vec;
}
 
int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
 
    // input
    string S;
    cin >> S;
    
    // calculation
    if ((int)S.find('W') != -1 && (int)S.find('B') != -1) {
        vector<ll> w_idx = findChar(S, 'W'), vec_Bnum(0);
        ll left = w_idx.front() - 0;  // 左端のBの数
        vec_Bnum.push_back(left);
        
        size_t len = w_idx.size();
        for (int i = 1; i < len; ++i) {
            vec_Bnum.push_back(w_idx[i] - w_idx[i-1] - 1);
        }
        
        partial_sum(vec_Bnum.begin(), vec_Bnum.end(), vec_Bnum.begin());
        if (vec_Bnum.size() == 0) {
            cout << 0 << "\n";
        } else {
            cout << SUM(vec_Bnum) << "\n";
        }
    } else {
        cout << 0 << "\n";
    }
    
    return 0;
}

#include <numeric>で使えるpartial_sum()は知れてよかった。

• partial_sum

このコードのWAの原因は探せなかった。計算手順的に間違ってないと思うんだけども...

解法

WAになった原因はよく分からないまま... vectorを使わずにシンプルに累積和を書いたらACできた。ソースコードは以下。

#include <iostream>
#include <cstdio>
#include <string>
using namespace std;
 
#define REP(i, n) for (int i=0; i < (int)(n); i++)  // 0 ~ n-1
typedef long long ll;
typedef unsigned long ul;
 
int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    
    // input
    string S;
    cin >> S;
    
    // calculation
    ll ans = 0, cnt = 0;
    ll len = (ll)S.length();
    REP(i, len) {
        if (S[i] == 'B') {
            cnt++;
        } else {
            ans += cnt;
        }
    }
  
    cout << ans << "\n";
    
    return 0;
}

こんなシンプルに累積和を書けるとは思わなかった。Bの数を保持しつつ、答えとなる和に加算していく方法で累積和を書けるのはなるほどと思った。