Epic Subarrays Counting (Editorial)

Solution

Notice that we cannot have an epic array in which $1$ appears anywhere except the first position , i.e. the general form of an epic array is

therefore we store positions of $1$'s into an array $\text{pos}$ to make the process easier.

For (1 \le i \le n-1) we get that the size of block of $2$'s between a_{\text{pos}[i]} and a_{\text{pos}[i+1]} is \text{pos}[i+1]-\text{pos}[i]-1 , since we're also counting the starting $1$ the expression becomes \text{pos}[i+1]-\text{pos}[i].

Now for each block it holds that we want to count the subarrays of the block itself , let starting index of $1$ by $a$ and ending index by $b$ then by simulating the process we end up with (a\rightarrow a) , (a \rightarrow a+1) , .. (a \rightarrow b-1) , similarly for (a+1 \rightarrow a+1) ,... , (a+1 \rightarrow b-1)

This simulation ends up being the same as the summation of first $k$ integers i.e. Triangular numbers T(n)=\frac{n(n+1)}{2}.

After that you've to consider the segment two that appears before a_{\text{pos}[0]} since it'll not be counted , as well as a_{n-\text{pos}[\text{last}]}.

and thus our solution is complete when $a$ contains at least single $2$.

When a=[1,1,1,...,1] the solution is $T(n)$.

Complexity $\mathcal{O}(n)$.

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;

#define int long long
template<class T> using oset=tree<T,null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;

int T(int x){return x*(x+1)/2;}

void solve(){
    int n;cin>>n;
    vector<int>a(n);
    for(auto &i:a)cin>>i;
    vector<int>pos;
    for(int i=0;i<n;i++){
        if(a[i]==1) pos.push_back(i);
    }
    int ans=0;
    if(pos.empty()){
        ans=T(n);
    }else{
        ans+=T(pos[0]);           
        for(int i=0;i+1<(int)pos.size();i++){
            ans+=T(pos[i+1]-pos[i]);   
        }
        ans+=T(n-pos.back());        
    }
    cout<<ans<<"\n";
}

signed main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int t;cin>>t;
    while(t--) solve();
    return 0;
}