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

ABC162 に(バーチャル)参加しました。
D問題が面白かったので初めての解説記事を書きます！

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

なお解いている様子は YouTube に動画があります。

問題

D - RGB Triplets

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

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

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 を加算する。

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 を連想した。

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

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本

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

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)$で求まる。

        # 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！

ソース全体

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)