【ABC162 D】「RGB Triplets」を解く【Python3】

「R」「G」「B」のみからなる、長さNの文字列Sに対して、以下の条件を共に満たす組 (i, j, k)を求める問題です。

「Si ≠ Sj」かつ「Si ≠ Sk」かつ「Sj ≠ Sk」
「j−i ≠ k−j」

提出

n = int(input())
s = input()

ans = s.count("R") * s.count("G") * s.count("B")

for i in range(n-2):
    for j in range(i+1, i+(n-i-1)//2+1):
        k = 2*j-i
        if k < n and s[k] != s[i] and s[k] != s[j] and s[i] != s[j]:
            ans -= 1

print(ans)

組(i, j, k)は必ず、「R」「G」「B」の3つを組み合わせたものになります。

文字列Sの中の「R」「G」「B」の数を数え、掛け合わせることで、すべてのパターンを導き出すことができます。

このパターンの中で、「j−i = k−j」であるものを除外するために、for文を使用し、

「j−i = k−j」で、「Si ≠ Sj」かつ「Si ≠ Sk」かつ「Sj ≠ Sk」のものが出てきたときには、答えに「-1」をします。

ここで、すべての(i, j, k)を試していたら、時間コストがかかります。

「j−i = k−j」から「k = 2*j-i」を導き出せるので、これを用いて、(i, j)のみを調べます。

さらにコストを減らすために、iを用いてjの範囲を決めます。