masuTomo’s blog

競技プログラミングの勉強メモです.主にPythonを使用します.数学の記事も書くかもしれません.

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重ループに減らすというものです.

おおよそのところを説明すると,整数Kが与えられたときに a+b+c+d = Kとなるa,b,c,dの組み合わせを求めよという問に対し,最初の4重ループはa,b,c,dをそれぞれforでまわして解を得るものでした.

でも,最後のdについてはKが決まっているのだから,a,b,cが決まった段階でdをループさせる必要はなくなるので3重ループに落とせるという考え方です.

これを真似します. 今の問題では左から順に調べていき,R,Gが使わていれば最後に探す文字はBになります.そこでGを使用した場所よりも右にあるBの個数がわかればそれで(i,j,k)の個数が確定します.

そのために,左から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)