fork download
  1. #include <bits/stdc++.h>
  2. #include <ext/pb_ds/assoc_container.hpp>
  3. #include <ext/pb_ds/tree_policy.hpp>
  4.  
  5. using namespace std;
  6. using namespace __gnu_pbds;
  7.  
  8. using ordered_set = tree<
  9. pair<int, int>,
  10. null_type,
  11. less<pair<int, int>>,
  12. rb_tree_tag,
  13. tree_order_statistics_node_update
  14. >;
  15.  
  16. int main() {
  17. int n;
  18. string s;
  19. cin >> n >> s;
  20.  
  21. ordered_set X;
  22.  
  23. X.insert({0, 0});
  24.  
  25. int prefix = 0;
  26. long long ans = 0;
  27.  
  28. for (int j = 0; j < n; j++) {
  29. prefix += (s[j] == '1' ? 1 : -1);
  30.  
  31. ans += X.order_of_key({prefix, 0});
  32.  
  33. X.insert({prefix, j + 1});
  34. }
  35.  
  36. cout << ans << '\n';
  37. }
Success #stdin #stdout 0s 5320KB
stdin
5
11010
stdout
8