D - RGB Triplets
・問題URL
https://atcoder.jp/contests/abc162/tasks/abc162_d
・発想
・ijk3つの全探索は無理
・ijだけ全探索すると間に合い、j+1番目からN番目までで、i,j番目の要素と異なるものの個数がわかればよさそう。これは3色とも累積和でO(N)で求められる
・2つ目の条件どーしよー
・解法
累積和でSのx文字目までのRGBの数を全部求めておき、i,jで全探索(O(N²))。
二つ目の条件を満たさないとき、iとkの平均がjゆえ、2*j-iがN以下ならijkが等間隔になりうるのでその時はk番目がかぶってないか調べて、かぶってたらansから1引く。
・コード
int main(){
int N;
string S;
cin>>N>>S;
vector<ll>r(N+1,0),g(N+1,0),b(N+1,0);
rep(i,N){
if(S[i]=='R'){
r[i+1]=r[i]+1;
g[i+1]=g[i];
b[i+1]=b[i];
}
else if(S[i]=='G'){
r[i+1]=r[i];
g[i+1]=g[i]+1;
b[i+1]=b[i];
}
else {
r[i+1]=r[i];
g[i+1]=g[i];
b[i+1]=b[i]+1;
}
}
ll ans=0;
for(int i=1;i<=N-2;i++){
for(int j=i+1;j<=N-1;j++){
if((S[i-1]=='R'&&S[j-1]=='G')||(S[i-1]=='G'&&S[j-1]=='R')){
ans+=(b[N]-b[j]);
if(2*j-i<=N){
if(S[2*j-i-1]=='B')ans--;
}
}
else if((S[i-1]=='G'&&S[j-1]=='B')||(S[i-1]=='B'&&S[j-1]=='G')){
ans+=(r[N]-r[j]);
if(2*j-i<=N){
if(S[2*j-i-1]=='R')ans--;
}
}
else if((S[i-1]=='B'&&S[j-1]=='R')||(S[i-1]=='R'&&S[j-1]=='B')){
ans+=(g[N]-g[j]);
if(2*j-i<=N){
if(S[2*j-i-1]=='G')ans--;
}
}
}
}
cout<<ans<<endl;
}
この記事が気に入ったらサポートをしてみませんか?