Atcoder - 162D 枚舉

<Atcoder - 162D> 枚舉

https://atcoder.jp/contests/abc162/tasks/abc162_d

題意：

給定一個只由字符 'B'，'G'，'R' 組成的一個串S，求滿足下列條件的三元組{Si，Sj，Sk}的個數：

(1) Si != Sj && Si != Sk && Sj != Sk

(2) j - i != k - j

思路：

先把字符串轉化成數字串(B =>1，G => 2，R => 3)，爲了方便操作。分別記B G R的個數爲cnt[1]，cnt[2]，cnt[3]，這樣能產生的滿足條件(1)的三元組個數共有cnt[1] * cnt[2] * cnt[3]個，然後我們寫一個chk函數，枚舉所有滿足條件(1)但是不滿足條件(2)的三元組，即看傳入的三元組是否相鄰的兩個元素間距相等。我們先將1，2，3對應的位置放進vector內，這樣不斷枚舉三元組內第一個元素的位置，然後第二層循環枚舉距離，在不越界串長度的條件下，看如果當前枚舉的首個數字的位置加上一倍距離和二倍距離是否分別是剩下的兩個數字。三元組一共就6種情況：123、132、213、231、312、321，分別用chk加工這6個串，然後用cnt[1] * cnt[2] * cnt[3]減去這些不合法的即可。

AC代碼：

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int maxx = 4e3 + 7;
int n;
char s[maxx];
int num[maxx];
vector <int> vec[maxx];
LL cnt[maxx];
bool vis[maxx];

LL chk(int a, int b, int c) {
	LL res = 0;
	for(int i = 0; i < vec[a].size(); i++) { //枚舉首字母位置
		for(int j = 1; vec[a][i] + 2 * j <= n; j++) { //枚舉間距
			if(num[vec[a][i] + j] == b && num[vec[a][i] + 2 * j] == c) res++;
		}
	}
	return res;
}

int main() {
	scanf("%d", &n);
	scanf("%s", s + 1);
	for(int i = 1; i <= n; i++) {
		if(s[i] == 'B') num[i] = 1;
		else if(s[i] == 'G') num[i] = 2;
		else if(s[i] == 'R') num[i] = 3;
	}
	for(int i = 1; i <= n; i++) {
		vec[num[i]].push_back(i);
		cnt[num[i]]++;
	}
	LL ans = cnt[1] * cnt[2] * cnt[3];
	ans -= chk(1, 2, 3) + chk(1, 3, 2) + chk(2, 1, 3) + chk(2, 3, 1) + chk(3, 1, 2) + chk(3, 2, 1);
	printf("%lld\n", ans);
	return 0;
}