AtCoder Beginner Contest 162 [ D - RGB Triplets ] https://atcoder.jp/contests/abc162/tasks/abc162_d

abc162_d.py

'''
[問題]
https://atcoder.jp/contests/abc162/tasks/abc162_d

[解法]
AtCoder Beginner Contest 143 [ D - Triangles ]をPythonで解く
https://uchidama.hatenablog.com/entry/2021/08/06/150713

この問題に似てる。
1 <= N <= 4000
なのでO(N^3)だと、4^3 * 10^9 でTLEのはず。

なので、R, Gの二つの要素を固定。
すると O(16 * 10^6) くらい？
あとは、計算でBの要素から何個採用できるかを出す。
この方針でTLEは出ないはず。

で、なんとかなった。

[結果]
PyPy3(7.3.0) AC 315ms
Python3.8.2 AC 1967ms   お、ギリやな
'''

import sys

sys.setrecursionlimit(10 ** 6)  # 再帰上限の引き上げ
input = sys.stdin.readline

INF = 2 ** 63 - 1

N = int(input())
S = input()
S = S.replace('\n', '')

R = set()
G = set()
B = set()

for i, s in enumerate(S):
    if s == 'R':
        R.add(i)
    elif s == 'G':
        G.add(i)
    else:
        B.add(i)

ans = 0

for r in R:
    for g in G:

        big = max(r, g)
        small = min(r, g)

        # Bの値が何個採用できるかを計算で出す。

        # big, smallが決まっているなら、big, smallがどの変数に当たるか仮定する
        # small = i, big = j
        # small = i, big = k
        # small = j, big = k
        # j-i != k-j である必要があるので、採用しない要素があるか確認する
        d_cnt = 0   # 削除するカウント数

        # small = i, big = j の時、kの場所にあるBは採用できない
        if (big - small + big) in B:
            d_cnt += 1
        # small = i, big = k の時、jの場所にあるBは採用できない
        if ((big - small)/2 + small) in B:
            d_cnt += 1
        # small = j, big = k の時、iの場所にあるBは採用できない
        if (small - (big - small)) in B:
            d_cnt += 1

        ans += max(0, len(B) - d_cnt)

print(ans)