Codeforces Round #780 F2. Promising String(hard version)

传送门

题目大意

一个长为 \(n(\leq n\leq 2\times 10^5)\)的由 \(‘+’\) 和 \(‘-’\) 组成的串，可以进行若干次操作，将相邻的两个 \('-'\) 替换为一个 \('+'\) 。求有多少个字串在执行若干次操作后（可以不操作），其中的 \('+'\) 与 \('-'\) 数量相等。

思路

我们可以发现只有当子串中 \('-'\) 的数量比 \('+'\) 的数量多出 \(3\) 的倍数个时才可以满足要求。我们可以求出 \('+'\) 与 \('-'\) 差的前缀和，然后枚举字串的右端点，用 \(3\) 个树状数组分别维护一下模 \(3\) 为 \(0,1,2\) 的前缀和的值的数量，由于会有负数，所以我们加上一个偏移 \(h\) ，并且一开始要有一个前缀和为 \(0\) 的值，之后对于每个右端点，查询有多少个比右端点处的前缀和值要小并且与其模 \(3\) 同余的值，加入到答案中即可，复杂度 \(O(nlogn)\) 。

代码

#include<bits/stdc++.h>
#include<unordered_map>
#include<unordered_set>
using namespace std;
using LL = long long;
using ULL = unsigned long long;
using PII = pair<int, int>;
using TP = tuple<int, int, int>;
#define all(x) x.begin(),x.end()
//#define int LL
//#define lc p*2
//#define rc p*2+1
#define endl '\n'
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f
//#pragma warning(disable : 4996)
#define IOS ios::sync_with_stdio(0),cin.tie(0),cout.tie(0)
const double eps = 1e-8;
const LL MOD = 1000000007;
const LL mod = 998244353;
const int maxn = 200010;

LL T, N;
LL D[maxn];
string S;
LL n, bit[3][maxn * 2];

void add(int i, int x, int m)
{
	while (i <= n)
	{
		bit[m][i] += x;
		i += i & (-i);
	}
}

LL sum(int i, int m)
{
	LL ans = 0;
	while (i)
	{
		ans += bit[m][i];
		i -= i & (-i);
	}

	return ans;
}

void solve()
{
	n = N * 2 + 1;
	LL ans = 0;
	LL h = N + 1;
	for (int i = 1; i <= N; i++)
		D[i] = D[i - 1] + (S[i - 1] == '+' ? -1 : 1);
	add(h, 1, 0);
	for (int i = 1; i <= N; i++)
	{
		int m = (D[i] % 3 + 3) % 3;
		LL num = sum(D[i] + h, m);
		ans += num;
		add(D[i] + h, 1, m);
	}
	cout << ans << endl;
}

int main()
{
	IOS;
	cin >> T;
	while (T--)
	{
		cin >> N >> S;
		for (int i = 1; i <= N; i++)
			D[i] = 0;
		for (int i = 1; i <= N * 2 + 1; i++)
			bit[0][i] = bit[1][i] = bit[2][i] = 0;
		solve();
	}

	return 0;
}