CF1660 F2. Promising String (hard version) （树状数组）

https://codeforces.com/contest/1660/problem/F2
题意:

思路： 设+个数为x， -个数为y， (y - x) % 3 == 0 && y - x >= 0 时为希望串。
+定义为-1，-定义为+1，做普通前缀和p和%3的前缀和pre。那么枚举左区间，右区间就是 pr和pl相同且prer - prel-1 的值大于零0的所有点。 权值线段树或树状数组解决。

#include<bits/stdc++.h>
using namespace std;
#define IOS ios::sync_with_stdio(false) ,cin.tie(0), cout.tie(0);
//#pragma GCC optimize(3,"Ofast","inline")
#define ll long long
#define PII pair<int, char>
//#define int long long
const int N = 4e5 + 5;
const int M = 1e5 + 5;
const int INF = 0x3f3f3f3f;
const ll LNF = 0x3f3f3f3f3f3f3f3f;
const int mod = 998244353;
const double PI = acos(-1.0);
int pre[N], p[N];
int tr[3][N];
int lowbit( int x ) { return x & (-x); } 
int n;
void add ( int x, int ty ) {
  for ( int i = x; i <= (n << 1) + 2; i += lowbit(i) ) tr[ty][i] += 1;
  
}
ll query ( int x, int ty) {
  ll res = 0;
  for ( int i = x; i >= 1; i -= lowbit(i) ) {
    res += tr[ty][i];
  }
  return res;
}
void solve() {
  cin >> n;
  string ss; cin >> ss; ss = "*" + ss;

  for ( int i = 1; i <= n; ++ i ) {
    if(ss[i] == '-') ++ pre[i], ++ p[i]; else -- pre[i], --p[i]; pre[i] += pre[i - 1], p[i] += p[i - 1];
    pre[i] = (pre[i] + 3) % 3;
    //cout << pre[i] << " " << p[i] << endl;
  }

  ll ans = 0;

  for ( int i = n; i >= 0; -- i ) {
    // int tt = pre[i]; int sb = tr[2][4], l =-p[i] + n + 2 - 1;
    ans += query( (n << 1) + 2, pre[i]) - query( p[i] + n + 2 - 1, pre[i]);
    add( p[i] + n + 2, pre[i]);
  }

  for ( int i = 1; i <= (n << 1) + 2; ++ i ) {
    pre[i] = 0; p[i] = 0;
    for ( int j = 0; j <= 2; ++ j ) tr[j][i] = 0;
  }
  cout << ans << '\n';
}
int main () {
  int t; cin >> t;
  while ( t -- ) solve();
  return 0;
}