(Python) ABC162-D 考察ログと解法

[ABC162](https://atcoder.jp/contests/abc162) に(バーチャル)参加しました。
D問題が面白かったので初めての解説記事を書きます！

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/155933/febfd5fc-ca05-0827-eb9e-e89b13037c28.png)

綺麗な解説ではなく、その時思ったことを垂れ流すように書いてみました。

なお解いている様子は [YouTube](https://www.youtube.com/watch?v=d2Ha71q17_c) に動画があります。

# 問題
[D - RGB Triplets](https://atcoder.jp/contests/abc162/tasks/abc162_d)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/155933/ca95e703-afd6-85c5-aca3-bd3c20b41970.png)

# 入力を受け取りながら問題文を理解する

脳死でやれることをやりながら問題を読む。まずは入力を受け取ろう。

```python3
n = int(input())
s = input()
```

文中の2つの条件は、つまりこういうことか。

* 文字列 $S$ から、 `R`, `G`, `B`を一つずつ取り出す方法は何通りか？
* ただし、 $i, j, k$ が等間隔に並ぶのはだめ。 $(i, j, k) = (1, 2, 3), (3, 5, 7)$ など。

# 愚直を書きながら考察する

困ったら愚直。手を動かしながら考えよう。

$i < j < k$となる組み合わせを全探索して、条件に一致したら `ans` を加算する。

```python3
ans = 0

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

print(ans)
```

これは明らかに $O(N^3)$ 。$N<=4000$ なのでだめか。
$O(N^2)$ にすれば間に合うけど、ループ回数減らせるかな？

# 2つ決めたら残り1つは自動で決まる

3重ループを見たら、 [Otoshidama](https://atcoder.jp/contests/abc085/tasks/abc085_c) を連想した。

2重ループにできるかな。こんな処理ができれば$O(N^2)$でいけるか？

```python3
for i in range(n-2):
    for j in range(i+1, n-1):
        if s[i] == s[j]:
            continue
        c = "(s[i],s[j]と異なる、(R,G,B)のどれか)"

        ans += "s[j+1:n] の中の c の個数"

        # ただし、i,j,kが等間隔に並ぶ位置は不適なので1減らす
        if j+(j-i)<n and s[j+(j-i)] == c:
            ans -= 1
```

# 累積和3本

「"区間中の個数" は累積和！(素振り)」ということで、累積和でしょう。
ここまでくればあとは累積和を作るだけ。

```python3
ruisekiR = [0] * (n+1)
ruisekiG = [0] * (n+1)
ruisekiB = [0] * (n+1)

for i,c in enumerate(s):
    ruisekiR[i+1] = ruisekiR[i] + ( 1 if c == "R" else 0)
    ruisekiG[i+1] = ruisekiG[i] + ( 1 if c == "G" else 0)
    ruisekiB[i+1] = ruisekiB[i] + ( 1 if c == "B" else 0)
```

これを作っておけば、 `s[j+1:n] の中の c の個数` は$O(1)$で求まる。

```python3
        # c = (s[i],s[j]と異なる、(R,G,B)のどれか)
        c = "R"
        if "R" in (s[i], s[j]):
            c = "G"
            if "G" in (s[i], s[j]):
                c = "B"

        # ans += (s[j+1:n] の中の c の個数)
        if c == "R":
            ans += (ruisekiR[n] - ruisekiR[j])
        if c == "G":
            ans += (ruisekiG[n] - ruisekiG[j])
        if c == "B":
            ans += (ruisekiB[n] - ruisekiB[j])
```

なんか汚いけどコンテストなので気にしない。
これでAC！

# ソース全体

```python3
n = int(input())
s = input()

ruisekiR = [0] * (n+1)
ruisekiG = [0] * (n+1)
ruisekiB = [0] * (n+1)

for i,c in enumerate(s):
    ruisekiR[i+1] = ruisekiR[i] + ( 1 if c == "R" else 0)
    ruisekiG[i+1] = ruisekiG[i] + ( 1 if c == "G" else 0)
    ruisekiB[i+1] = ruisekiB[i] + ( 1 if c == "B" else 0)

ans = 0

for i in range(n-2):
    for j in range(i+1, n-1):
        if s[i] == s[j]:
            continue
        c = "R"
        if "R" in (s[i], s[j]):
            c = "G"
            if "G" in (s[i], s[j]):
                c = "B"

        if c == "R":
            ans += (ruisekiR[n] - ruisekiR[j])
        if c == "G":
            ans += (ruisekiG[n] - ruisekiG[j])
        if c == "B":
            ans += (ruisekiB[n] - ruisekiB[j])

        if j+(j-i)<n and s[j+(j-i)] == c:
            ans -= 1

print(ans)
```