Обчислення епічних підмасивів (Розбір)

Рішення

Зверніть увагу, що ми не можемо мати епічний масив, в якому $1$ з'являється де-небудь, крім першої позиції, тобто загальна форма епічного масиву є

тому ми зберігаємо позиції $1$ в масиві $\text{pos}$, щоб спростити процес.

Для (1 \le i \le n-1) ми отримуємо, що розмір блоку $2$ між a_{\text{pos}[i]} і a_{\text{pos}[i+1]} дорівнює \text{pos}[i+1]-\text{pos}[i]-1, оскільки ми також враховуємо початковий $1$, вираз стає \text{pos}[i+1]-\text{pos}[i].

Тепер для кожного блоку справедливо, що ми хочемо підрахувати підмасиви самого блоку, нехай початковий індекс $1$ буде $a$, а кінцевий індекс $b$, тоді, симулюючи процес, ми отримуємо (a\rightarrow a) , (a \rightarrow a+1) , .. (a \rightarrow b-1), аналогічно для (a+1 \rightarrow a+1) ,... , (a+1 \rightarrow b-1).

Ця симуляція виявляється такою ж, як сума перших $k$ цілих чисел, тобто трикутні числа T(n)=\frac{n(n+1)}{2}.

Після цього вам потрібно врахувати сегмент два, який з'являється перед a_{\text{pos}[0]}, оскільки він не буде врахований, а також a_{n-\text{pos}[\text{last}]}.

Таким чином, наше рішення завершено, коли $a$ містить принаймні один $2$.

Коли a=[1,1,1,...,1], рішення дорівнює $T(n)$.

Складність $\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;
}