Codeforces Round #833 (Div. 2) C. Zero-Sum Prefixes

Codeforces Round #833 (Div. 2) C. Zero-Sum Prefixes

Let’s consider the prefix sum array $s=[a_1,a_1+a_2,\dots ,a_1+a_2+\dots +a_n]$.

For every index $i$ such that $a_i=0$, if we change the value of $a_i$ to $x$, then every element from the suffix $[s_i,s_{i+1},\dots ,s_n]$ will be increased by $x$. Therefore, if $a_{i_1}=a_{i_2}=\dots =a_{i_k}=0$, we’ll partition array $s$ into multiple subarrays:

• $[s_1,s_2,\dots ,s_{i_1-1}]$;
• $[s_{i_2},s_{i_2+1},\dots ,s_{i_3-1}]$;
  $\dots$
• $[s_{i_k},s_{i_{k+1}},\dots ,s_n]$;
  Since none of the elements from the first subarray can be changed, it will contribute with the number of occurences of $0$ in $[s_1,s_2,\dots ,s_{i_1-1}]$ towards the final answer.

For each of the other subarrays $[s_l,s_{l+1},\dots ,s_r]$, let $x$ be the most frequent element in the subarray, appearing $fr[x]$ times. Since $a_l=0$, we can change the value of $a_l$ to $-x$. In this case, every $x$ in this subarray will become equal to $0$, and our current subarray will contribute with $fr[x]$ towards the final answer.

Time complexity per testcase: $O(NlogN)$

#include<bits/stdc++.h>
using namespace std;

typedef long long LL;

int main()
{
    int T;cin>>T;
    while(T--)
    {
        int n;cin>>n;

        map<LL,LL> b;
        bool ok_zero=false;
        LL sum=0,mx=0,res=0;

        for(int i=0;i<n;i++)
        {
            int x;cin>>x;

            if(x==0)
            {
                if(ok_zero) res+=mx;
                else res+=b[0];

                ok_zero=true;
                mx=0;
                b.clear();
            }

            sum+=x;
            b[sum]++;
            mx=max(mx,b[sum]);
        }

        if(ok_zero) res+=mx;
        else res+=b[0];

        cout<<res<<endl;
    }
    return 0;
}