[AtCoder] ABC 162 D – RGB Triplets

2020年4月13日2020年12月15日

問題

方針

まず初めに、\( S_i \neq S_j \wedge S_i \neq S_k \wedge S_j \neq S_j \) を満たす組の数を求めます。これは、\( i, j \) について、\( S_i \neq S_j \) となるものを求め、\( j + 1 \) 文字目以降に存在する \( S_i, S_j \) 以外の文字数を数えます。これは、累積和を使うことで高速に求まります。

つぎに、\( j – i = k – j \) を満たす \( i, j ,k \) について、\( S_i \neq S_j \wedge S_i \neq S_k \wedge S_j \neq S_j \) を満たすものを数えます。これは、\( i, j \) について全探索を行います。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int main() {
    int N;
    string S;
    cin >> N >> S;
    ll a[N + 1][3]{};
    map<char, int> m;
    m['R'] = 0;
    m['G'] = 1;
    m['B'] = 2;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < 3; j++) a[i + 1][j] = a[i][j];
        a[i + 1][m[S[i]]]++;
    }
    ll ans = 0;
    for (int i = 0; i < N - 2; i++) {
        for (int j = i + 1; j < N - 1; j++) {
            if (S[i] != S[j]) {
                int id = 0;
                for (int k = 0; k < 3; k++) {
                    if (k != m[S[i]] && k != m[S[j]]) {
                        id = k;
                        break;
                    }
                }
                ans += a[N][id] - a[j + 1][id];
            }
        }
    }
    for (int i = 0; i < N - 2; i++) {
        for (int j = i + 1; j < N - 1; j++) {
            int k = 2 * j - i;
            if (k < N && j < k && k >= 0) {
                if (S[i] != S[j] && S[j] != S[k] && S[i] != S[k]) ans--;
            }
        }
    }
    cout << ans << "\n";
    return 0;
}