Submission #72507930

---

Source Code Expand

#include <bits/stdc++.h>
 
using namespace std;
#define int long long
#define all(x) x.begin(), x.end()
 
#ifndef ONLINE_JUDGE
#include "debug.h"
#define debug(x...) cerr << "[" << #x << "] = ["; _print(x)
#else
#define debug(x...)
#endif
 
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
 
template <class T> void output(vector < T > & a) { for (auto & x: a) {cout << x << ' ';}cout << '\n';}


 
// Golden Rules
/*
    Solutions are simple.
 
    Proofs are simple.
 
    Implementations are simple.
 
                         ~ TEJA DRONADULA SIR!!
*/

// source : cp-algorithm
struct FenwickTree {
vector<int> bit;  // binary indexed tree
int n;

FenwickTree(int n) {
     this->n = n;
     bit.assign(n, 0);
}

FenwickTree(vector<int> const &a) : FenwickTree(a.size()) {
     for (size_t i = 0; i < a.size(); i++)
          add(i, a[i]);
}

int sum(int r) {
     int ret = 0;
     for (; r >= 0; r = (r & (r + 1)) - 1)
          ret += bit[r];
     return ret;
}

int sum(int l, int r) {
     return sum(r) - sum(l - 1);
}

void add(int idx, int delta) {
     for (; idx < n; idx = idx | (idx + 1))
          bit[idx] += delta;
}
};

void solve() {
     int N;
     cin >> N;
     string S;
     cin >> S;
     vector<int> pre(N + 1);
     for (int i = 0; i < N; ++i) {
          pre[i + 1] = (S[i] == 'A' ? 1 : S[i] == 'B' ? -1 : 0) + pre[i];
     }
     for (int i = 0; i <= N; ++i) {
          pre[i] += N;
     }
     FenwickTree fen(2 * N + 1);
     int ans = 0;
     for (int i = 0; i <= N; ++i) {
          ans += fen.sum(pre[i] - 1);
          fen.add(pre[i], 1);
     }
     cout << ans << "\n";
}


int32_t main() {
     ios::sync_with_stdio(false);
     cin.tie(nullptr);
 
     #ifndef ONLINE_JUDGE
          freopen("input.txt", "r", stdin);
          freopen("output.txt", "w", stdout);
          freopen("error.txt", "w", stderr);
     #endif

     int t = 1;
     // cin >> t; 

     for (int cases = 1; cases <= t; ++cases) {
          solve();
     }
 
     #ifndef ONLINE_JUDGE
          double T = 1000.0 * clock() / CLOCKS_PER_SEC;
          cout << "\n[Finished in " << T << "ms]";
          cerr << "\n[Finished in " << T << "ms]";
     #endif
 
     return 0;
}

Submission Info

Judge Result