けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 162 D - RGB Triplets (400 点)

落ち着いて頭を整理。。。

問題へのリンク

問題概要

'R', 'G', 'B' のみからなる長さ  N の文字列  S が与えられる。 S の index の組  1 \le i \lt j \lt k \le N であって、

  •  S_{i} S_{j} S_{j} はすべて互いに異なる
  •  j - i \neq k - j である

という条件を満たすものの個数を求めよ。

制約

  •  1 \le N \le 4000

考えたこと

頭がごっちゃになりそうだ。こういうのはきっちりと落ち着いて整理して数え上げることが求められる。とりあえず

  •  j - i \neq k - j である

という条件はとっても扱いづらい!!!!!この条件は「3 つの index の間隔が等しく無い」ということを言っているわけだけど、むしろ「3 つの index の間隔が等しい」という方を数える方が簡単だ。というわけで大まかな方針としては

  • 全体から
  •  j - i = k - j であるケースを引く

という風に考えることにしよう。もう一つの「S[i], S[j], S[k] が異なる」という条件を加味すると、結局、次のような方針で問題を解くことができる。


  • 「S[i], S[j], S[k] が異なる」という条件を満たすものすべてから、
  • 「S[i], S[j], S[k] が異なる」かつ「j - i = k - j である」という条件を満たすものを引く

S[i], S[j], S[k] が異なるもの

これは要するに S の文字のうち R が r 個、G が g 個、B が b 個あるとしたら

  • r 個の R から 1 個選び
  • g 個の G から 1 個選び
  • b 個の B から 1 個選ぶ

という方法の数に対応する。つまり、r × g × b 通りだ。

「S[i], S[j], S[k] が異なる」かつ「j - i = k - j である」という条件を満たすもの

i, j, k が等差数列にならなければならないという制約がついているので、効率よく数えることができる。i と j を固定してあげると、k が自動的に決まるのだ。よって整理すると、


  • i, j (i < j) を固定すると
  • k = 2j - i と決まるので (この時点で k >= N の場合はダメ)
  • S[i], S[j], S[k] が互いに異なるならばカウントする

という風にすれば OK。計算量は  O(N^{2}) でできる。

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N; string S;
    cin >> N >> S;
    long long r = 0, g = 0, b = 0;
    for (auto c : S) {
        if (c == 'R') ++r;
        else if (c == 'G') ++g;
        else ++b;
    }
    long long all = r * g * b;
    long long sub = 0;
    for (int i = 0; i < N; ++i) {
        for (int j = i+1; j < N; ++j) {
            if (S[i] == S[j]) continue;
            int k = j*2-i;
            if (k >= N || S[k] == S[i] || S[k] == S[j]) continue;
            ++sub;
        }
    }
    cout << all - sub << endl;
}
 ```