CF1400-D. Zigzags

CF1400-D. Zigzags

题意：

给出一个由\(n\)个数字构成的数组\(a\)，让你在这个数组中找出有多少个符合以下要求的元组\((i,j,k,l)\)：

1. \(i<j<k<l\)；

2. \(a_i==a_k,a_j=a_l\).

思路：

维护两个前缀和\(pre,suf\)。\(pre\)维护前\(i\)个数字中数字\(j\)的数量，\(suf\)维护后\(i\)个数字中数字\(j\)的数量。

那么只需要枚举每一个\(j\)和\(k\)，用前缀和找出前\(j-1\)个数字中\(a[k]\)的数量\(pre[j-1][a[k]]\)和后\(k+1\)个数字中\(a[j]\)的数量\(suf[k+1][a[j]]\)，两个数字相乘就是对于当前\(j,k\)符合要求的元组的数量。最后累加起来就是答案。

\(ans=\sum_{2<=j<k<=n-1}pre[j-1][a[k]]*suf[k+1][a[j]]\)

---

AC代码：

代码中的下标都是从\(0\)开始的，思路中提到的下标都是从\(1\)开始的。

#include <cstdio>
#include <cstring>
#include <iostream>

typedef long long ll;

const int Maxn = 3005;

int pre[Maxn][Maxn], suf[Maxn][Maxn], a[Maxn];

void solve() {
	memset(pre, 0, sizeof pre);
	memset(suf, 0, sizeof suf);
	int n;
	scanf("%d", &n);
	for (int i = 0; i < n; i++) {
		scanf("%d", a + i);
	}
	pre[0][a[0]]++;
	for (int i = 1; i < n; i++) {
		for (int j = 0; j <= n; j++) {
			pre[i][j] = pre[i - 1][j];
		}
		pre[i][a[i]]++;
	}
	suf[n - 1][a[n - 1]]++;
	for (int i = n - 2; i >= 0; i--) {
		for (int j = 0; j <= n; j++) {
			suf[i][j] = suf[i + 1][j];
		}
		suf[i][a[i]]++;
	}
	ll ans = 0;
	for (int j = 1; j <= n - 3; j++) {
		for (int k = j + 1; k <= n - 2; k++) {
			ans += 1LL * pre[j - 1][a[k]] * suf[k + 1][a[j]];
		}
	}
	printf("%lld\n", ans);
}

int main() {
	int T;
	scanf("%d", &T);
	while (T--) {
		solve();
	}
	return 0;
}