047 - Monochromatic Diagonal（★7）- Z-algorithm (ACL) を用いた解法。(naoya_t)

047.cc

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
 
typedef vector<int> vi;
#define pb  push_back
#define rep(var,n)  for(int var=0;var<(n);++var)
#define ALL(c)  (c).begin(),(c).end()
 
int mp[256];
 
int solve(int N, string& s, string& t) {
    int ans = 0;
    vi a(N * 2);
 
    rep(i,N){
        a[i] = (3 - mp[t[i]]) % 3;
        a[N + i] = mp[s[i]];
    }
    rep(t, 3) {
        vi z = atcoder::z_algorithm(a);
        rep(i, N){
            if (z[N+i] == N-i) ++ans;
        }
        rep(i,N){
            a[i] = (a[i] + 1) % 3;
        }
    }
 
    reverse(ALL(s));
    reverse(ALL(t));
 
    rep(i,N){
        a[i] = (3 - mp[t[i]]) % 3;
        a[N + i] = mp[s[i]];
    }
    rep(t, 3) {
        vi z = atcoder::z_algorithm(a);
        rep(i, N){
            if (i == 0) continue;
            if (z[N+i] == N-i) ++ans;
        }
        rep(i,N){
            a[i] = (a[i] + 1) % 3;
        }
    }
    return ans;
}
 
int main() {
    mp['R'] = 0;
    mp['G'] = 1;
    mp['B'] = 2;
 
    int N; cin >> N;
    string s, t; cin >> s >> t;
    cout << solve(N,s,t) << endl;
    return 0;
}