fork download
  1. #include <iostream>
  2. #include <bits/stdc++.h>
  3. using namespace std;
  4.  
  5. long long countStableSubarrays(vector<int>& capacity) {
  6. int n = capacity.size();
  7.  
  8. vector<long long> prefix(n, 0);
  9. prefix[0] = capacity[0];
  10.  
  11. for (int i = 1; i < n; i++) {
  12. prefix[i] = prefix[i - 1] + capacity[i];
  13. }
  14.  
  15. map<pair<long long, long long>, long long> mpp;
  16.  
  17. long long ans = 0;
  18.  
  19. for (int i = 0; i < n; i++) {
  20. long long target = prefix[i] - (long long)2 * (long long)capacity[i];
  21.  
  22. if (mpp.find({target, capacity[i]}) != mpp.end()) {
  23. ans = ans + mpp[{target, capacity[i]}];
  24. }
  25.  
  26. if (i != 0 && capacity[i] == 0 && capacity[i] == capacity[i - 1]) {
  27. ans = ans - 1;
  28. }
  29.  
  30. mpp[{prefix[i], capacity[i]}]++;
  31. }
  32.  
  33. return ans;
  34. }
  35.  
  36. int main() {
  37. int n;
  38. cin >> n;
  39.  
  40. vector<int> capacity(n);
  41.  
  42. for (int i = 0; i < n; i++) {
  43. cin >> capacity[i];
  44. }
  45.  
  46. cout << countStableSubarrays(capacity);
  47.  
  48. return 0;
  49. }
Success #stdin #stdout 0.01s 5288KB
stdin
5
9 3 3 3 9
stdout
2