[CF1400D] Zigzags

[CF1400D] Zigzags - 二分

Description

给定一个长度为 (n le 3000) 的数列，求有多少个四元组 ((i,j,k,l)) 满足 (1 le i < j < k < l le n) 且 (a_i = a_k, a_j = a_l)。

Solution

枚举 (j,k)，用前缀和统计 (i)，用后缀和统计 (l)。

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

#define int long long
const int N = 1e+6 + 5;
const int M = 1e+3 + 5;
const int mod1 = 1e+9 + 7;
const int mod2 = 998244353;

#define dbg(x) cerr << #x << ":" << x << endl

void solve()
{
	int n;

	cin >> n;

	vector<int> a(n + 2, 0);
	for (int i = 1; i <= n; i++)
	{
		cin >> a[i];
	}
	vector<vector<int>> b(n + 2);

	for (int i = 1; i <= n; i++)
	{
		b[a[i]].push_back(i);
	}

	auto countLess = [&](int x, int pos) -> int {
		auto t = lower_bound(b[x].begin(), b[x].end(), pos);
		return t - b[x].begin();
	};

	auto countMore = [&](int x, int pos) -> int {
		auto t = upper_bound(b[x].begin(), b[x].end(), pos);
		return b[x].end() - t;
	};

	int ans = 0;

	for (int j = 1; j <= n; j++)
	{
		for (int k = j + 1; k <= n; k++)
		{
			ans += countLess(a[k], j) * countMore(a[j], k);
		}
	}

	cout << ans << endl;
}

signed main()
{
	ios::sync_with_stdio(false);

	int t;
	cin >> t;
	while (t--)
	{
		solve();
	}

	return 0;
}