Codeforces #632 div.2 problem C

Can somebody help me solve this problem? I can calculate the subarrays with sum =0 but I am unable to count the good ones as the question requires because many subarrays overlap.

Any help would be appreciated, Thanks?

1 Like
//https://codeforces.com/contest/1333/problem/C
#include <iostream>
#include <bits/stdc++.h>
#include <cmath>
#include <vector>
#define ll long long int
#define mp make_pair
#define pb push_back
#define vi vector<int>
using namespace std;
ll solve(){
    int n;
    cin>>n;
    vector<int> seq(n);
    set<pair<ll,int>> prefixsums;
    ll sum=0;
    prefixsums.insert(mp(0,0));
    for(int i=0;i<n;i++){
        cin>>seq[i];
        sum+=seq[i];
        prefixsums.insert(mp(sum,i+1));
//keep a set with all prefix sums, and their index
    }
    set<pair<int,int>> badarray;
    //stores the first subarray with sum 0 from each starting index sorted by ending index
    vector<int> badarrayfromhere(n+1);
    //stores the first index at which the sum is 0, or equivalently, prefix sums are equal 
    for(auto x : prefixsums){
        int sidx,eidx;
        sidx=x.second;
        auto y=*(prefixsums.upper_bound(x));
        if(y.first!=x.first){//if there is no subarray with 0 sum from here 
            eidx=n+1;
            badarray.insert(mp(n+1,sidx));
        }
        else{
            eidx=y.second;
            badarray.insert(mp(eidx,sidx));//storing it as mentioned earlier.
        }
        badarrayfromhere[sidx]=eidx;//This is to remove them efficiently from the set
    }
    ll ans=0;
    for(int i=0;i<n;i++){
        The first point at which there will be a subarray of size 0.
        auto x=*(badarray.begin());
        ans+=x.first - i - 1;//The most subarrays possible from here
        badarray.erase(mp(badarrayfromhere[i],i));
        //We don't need to care about an element being equal to this anymore 
    }
    cout<<ans<<'\n';
}
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    solve();
}

It seemed like a straightforward question of sliding window but no matter what I did I wasn’t getting it right. Just waiting for editorial!