AtCoder Beginner Contest 162 D - RGB Triplets【Python】

https://atcoder.jp/contests/abc162/tasks/abc162_d

AtCoder ProblemsでDifficulty: 757、Solve Probability: 33%でした。

先にR,G,Bそれぞれの数の累積和を求めて、i, jについて全探索してkについては累積和から条件に合う数を求めました。
解説を見たところ、R,G,Bそれぞれの数をかけたものから２つ目の条件のものを引けばよかったようです。

N = int(input())
S = input()

sum_R = [0]
sum_G = [0]
sum_B = [0]
for i in range(N):
    if S[i] == "R":
        sum_R.append(sum_R[i] + 1)
        sum_G.append(sum_G[i])
        sum_B.append(sum_B[i])
    elif S[i] == "G":
        sum_R.append(sum_R[i])
        sum_G.append(sum_G[i] + 1)
        sum_B.append(sum_B[i])
    else:
        sum_R.append(sum_R[i])
        sum_G.append(sum_G[i])
        sum_B.append(sum_B[i] + 1)

ans = 0
for i in range(N):
    for j in range(i+1, N):
        s1 = S[i]
        s2 = S[j]
        ji_diff = j - i
        if (s1 == "R" and s2 =="G") or (s1 == "G" and s2 =="R"):
            ans += sum_B[N] - sum_B[j+1]
            # ２つ目の条件を満たすようにする
            if j + ji_diff < N:
                if S[j+ji_diff] == "B":
                    ans -= 1
        elif (s1 == "R" and s2 =="B") or (s1 == "B" and s2 =="R"):
            ans += sum_G[N] - sum_G[j+1]
            if j + ji_diff < N:
                if S[j+ji_diff] == "G":
                    ans -= 1
        elif (s1 == "G" and s2 =="B") or (s1 == "B" and s2 =="G"):
            ans += sum_R[N] - sum_R[j+1]
            if j + ji_diff < N:
                if S[j+ji_diff] == "R":
                    ans -= 1
print(ans)

解説を見ての実装

N = int(input())
S = input()

num_R = S.count("R")
num_G = S.count("G")
num_B = S.count("B")
All = num_R * num_G * num_B

ans = All
for i in range(N):
    for j in range(i+1, N):
        s1 = S[i]
        s2 = S[j]
        if j + (j - i) >= N:
            continue
        k = j + (j - i)
        s3 = S[k]
        if s1 == s2:
            continue
        if s1 == s3:
            continue
        if s2 == s3:
            continue
        ans -= 1
print(ans)