ABC162 解説(D)
ABC162 D - RGB Triplets
(ABC162 D - RGB Triplets) https://atcoder.jp/contests/abc162/tasks/abc162_d
TLEした解法
とりあえず何も気にせずに3重ループを書いて提出した.
N = int(input()) S = input() ans = 0 for i in range(N): for j in range(i,N): for k in range(j,N): if j -i == k-j : continue if S[i] != S[j] and S[j] != S[k] and S[k] != S[i]: ans += 1 print(ans)
結果:7TLE
制約はN<=4000なので,まぁそれはそうかと思っていましたが一応提出しました.
やり直し
3重ループをなんとかして2重ループにすれば通るだろうなと考えたので,その方法を考察します.
なんとなく頭の片隅にあったのは蟻本のp.8「あみだくじ」やp.25「ハードルがあがった『くじびき』」でした.これは最初4重ループで記載していたコードを3重ループに減らすというものです.
おおよそのところを説明すると,整数が与えられたときにとなるの組み合わせを求めよという問に対し,最初の4重ループはをそれぞれforでまわして解を得るものでした.
でも,最後のについてはが決まっているのだから,が決まった段階でをループさせる必要はなくなるので3重ループに落とせるという考え方です.
これを真似します. 今の問題では左から順に調べていき,R,Gが使わていれば最後に探す文字はBになります.そこでGを使用した場所よりも右にあるBの個数がわかればそれでの個数が確定します.
そのために,左からi番目までにR,G,Bの出現した個数を累積和としてあらかじめ計算しておくことにしました. あとはR,Gについて,2重ループをして,最後のBについては累積和から個数を求めました. コードが冗長になりましたが,ACしたものはこちら(PyPyで提出)
N = int(input()) S = input() #i番目までに何個のR,G,Bが出てくるか countR = [0 for _ in range(N)] countG = [0 for _ in range(N)] countB = [0 for _ in range(N)] for i in range(N): if S[i] == 'R': countR[i] = 1 if S[i] == 'G': countG[i] = 1 if S[i] == 'B': countB[i] = 1 for i in range(N-1): countR[i+1] += countR[i] countG[i+1] += countG[i] countB[i+1] += countB[i] ans = 0 #フラグ R = False G = False B = False for i in range(N): if S[i] == 'R': R = True if S[i] == 'G': G = True if S[i] == 'B': B = True for j in range(i,N): if S[i] == S[j]: continue if S[j] == 'R': if G: ans += countB[N-1] - countB[j] if 2*j-i < N and S[2*j-i] == 'B':ans -=1 if B: ans += countG[N-1] - countG[j] if 2*j-i < N and S[2*j-i] == 'G':ans -=1 if S[j] == 'G': if R: ans += countB[N-1] - countB[j] if 2*j-i < N and S[2*j-i] == 'B':ans -=1 if B: ans += countR[N-1] - countR[j] if 2*j-i < N and S[2*j-i] == 'R':ans -=1 if S[j] == 'B': if R: ans += countG[N-1] - countG[j] if 2*j-i < N and S[2*j-i] == 'G':ans -=1 if G: ans += countR[N-1] - countR[j] if 2*j-i < N and S[2*j-i] == 'R':ans -=1 R = False G = False B = False print(ans)