AtCoder Beginner Contest 162 D ( Python (3.8.2))

Atcoder Beginner Contests 162 D
競プロを本格的に始めて３ヶ月ちょっと経ちました。
自分なりに解いた過去問を解説していこうと思います。

※既にもっと素晴らしい記事を書かれている方がたくさんおり、悪魔で自分のメモがてらとして投稿しているのでご理解の程をよろしくお願い致します。

atcoder.jp

まず初めに考えたこと

制約が 1 <= N <= 4000 なのでfor３重for文で全探索しようとすると、O(N^3)になってしまい間に合わないことに気づく。

1つ目の条件

Si ≠ SjかつSi ≠ SkかつSj ≠ Skである
とあるが、これはN文字の中にある"R", "G", "B"の文字数を掛け合わせればいいだけ。

2つ目の条件

j−i ≠ k−jである
という条件指定があるので1つ目の条件から、この条件に当てはまる解を引くことで答えが求まる。

また、この条件からi, j を固定することでk = 2j - iと求まる。

kがこのままであると範囲を超えてしまうので、

0 <= 2j - i < n の制約をつけて調べれば終わり。

提出コードは以下の通り

n = int(input())

s = input()

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

cnt = 0 #2つ目の条件を満たさないもの

for i in range(n):

  for j in range(i+1,n):

    if 0 <= 2*j -i < n:

      if s[i] != s[j] and s[j] != s[2*j-i] and s[2*j-i] != s[i]:

        cnt += 1

print(num - cnt)

間違いなどがあれば指摘の方をよろしくお願いします。