Codeforces 780 F2. Promising String (hard version) (树状数组)

Codeforces 780 F2. Promising String (hard version) (树状数组)

题目

题目大意是一个字符串，每项为- 或者 +，如果-，+个数一样这个字符串是平衡的。你能进行操作：将两个相邻的-转变成+，如果一个字符串能通过操作变成平衡的被称为 “有希望平衡的” 字符串。问在给出的字符串的所有子串中有多少是“有希望平衡”的子串，输出数量。

思路

easy version 暴力就能做。对每段子串，sum表示+数量，则len - sum是-数量，只要-数量大于+数量，且差值是3的倍数，则这段是有希望平衡的。暴力算一下就行。
在hard version中 n <= 2e5, 没法暴力算。我们发现我们要找一段子串[l...r] ,令 val = (r - l + 1) - 2 * (sum[r] - sum[l - 1]) ，只要val > 0 && val % 3 = 0则对答案有贡献。

这道题解法利用一个结论：如果区间[1,i]%3=a， 并且区间[1,j]%3=a (j < i) 则 区间[j, i] % 3 = 0, 这样我们只需要遍历一次1-i区间，求出结果[1,i]%3=a, 然后我们只要能快速计算出前i个前缀里面有几个结果是a。这里说的不太清楚，只是一个大概的hint。具体见代码，总之就是这个快速求得方法能用树状数组解决。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<map>
#include<queue>
#include<stack>
#include<vector>
#include<string>
#include<set>
#include<fstream>
using namespace std;
#define rep(i, a, n) for(int i = a; i <= n; ++ i)
#define per(i, a, n) for(int i = n; i >= a; -- i)
typedef long long ll;
typedef pair<int, int> PII;
typedef pair<ll, int> PLI;
typedef pair<ll, ll> PLL;
const int N = 2e6 + 50;
const int mod = 998244353;
const double Pi = acos(- 1.0);
const ll INF = 1e12;
const int G = 3, Gi = 332748118;
ll qpow(ll a, ll b, ll p) { ll res = 1; while(b){ if(b & 1) res = (res * a) % p; a = (a * a) % p; b >>= 1;} return res % p; }
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }

char s[N];
int n;
ll c[3][N];
ll lowbit(int x){ return x & -x; }

ll ask(ll i, ll x){
	int ans = 0;
	for(; x; x -= lowbit(x)) ans += c[i][x];
	return ans;
}

void add(ll i, ll x,ll y){
	for( ; x<=N ; x += lowbit(x)) c[i][x] += y;
}

void solve(){
    scanf("%d",&n);
    ll n2 = (n + 1) << 1;
    for(int i = 0; i <= n2; ++ i) c[0][i] = c[1][i] = c[2][i] = 0;
    scanf("%s",s + 1);

    ll sum = 0, res = 0;
    add((n + 1) % 3, n + 1, 1);
    for(int i = 1; i <= n; ++ i){
        sum += (s[i] == '+');
        int tp = i - 2 * sum + n + 1;
        res += ask(tp % 3, tp);
        add(tp % 3, tp, 1);
    }

    printf("%lld\n",res);
}

int main() {
    freopen("temp.in", "r", stdin);
    freopen("temp.out", "w", stdout);
    int T;scanf("%d",&T);
    while(T --){
        solve();
    }
    return 0;
}