ABC162 D - RGB Triplets

備忘録

問題

atcoder.jp

回答

import sys
import os
import math

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

MOD = 10 ** 9 + 7


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N = ii()
    S = iss()
    mx = S.count('R') * S.count('G') * S.count('B')
    for i in range(N):
        for j in range(i + 1, N):
            k = j + j - i
            if k < N:
                if S[i] != S[j] and S[i] != S[k] and S[j] != S[k]:
                    mx -= 1

    print(mx)


if __name__ == '__main__':
    main()

考え方

まず、条件1(3つの文字が全て異なる)に当てはまるパターンの総数を考える。
対象の文字はR, G, Bのため、条件1はRの総数 * Gの総数 * Bの総数となる。
（文字列SがRRGBの場合、3つの文字が全て異なるパターンは2 * 1 * 1で2通り）

次に条件2(j - i ≠ k - j)について、
式をk = j + j - iに変形する。
k = j + j - iは条件2に当てはまらないパターンの回答、つまり
条件1の総数から条件2に当てはまらないパターンの総数をひくことで、
条件2に当てはまる（回答）を得ることができる。

実装は簡単で、
iとjからN以下となるkを求め、
条件2に当てはまる場合には、条件1の総数から1ひくことで、
回答を得ることが出来る。