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;
}