Codeforces Round #783 (Div. 1) B. Optimal Partition

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N=5e5+10;
const int mod= 998244353;
struct node{
    int l,r,val;
}tr[3][N<<2];
int t,n,a[N],s[N],now=0;
const int M=-1e9;
void pushup(int p){
    tr[now][p].val=max(tr[now][p<<1].val,tr[now][p<<1|1].val);
}
void build(int p,int l,int r){
    tr[now][p]={l,r,M};
    if(l==r) return ;
    int mid=(l+r)>>1;
    build(p<<1,l,mid);
    build(p<<1|1,mid+1,r);
    pushup(p);
}
void change(int pos,int p,int val){
    if(tr[now][p].l==tr[now][p].r){tr[now][p].val=max(val,tr[now][p].val);return ;}
    int mid=tr[now][p].l+tr[now][p].r>>1;
    if(pos<=mid) change(pos,p<<1,val);
    else change(pos,p<<1|1,val);
    pushup(p);
}
int query(int p,int l,int r){
    if(tr[now][p].l>=l&&tr[now][p].r<=r){
        return tr[now][p].val;
    }
    int res=-1e9;
    int mid=tr[now][p].l+tr[now][p].r>>1;
    if(mid>=l) res=max(res,query(p<<1,l,r));
    if(mid+1<=r) res=max(res,query(p<<1|1,l,r));
    return res;
}
map<int,int> mp;
int dp[N];
int32_t main(){
    scanf("%lld",&t);
    while(t--){
        vector<int> v;
        scanf("%d",&n);
        for(int i=1;i<=n;i++){
            scanf("%lld",&a[i]),s[i]=s[i-1]+a[i],v.push_back(s[i]);
            dp[i]=-1e9; 
        }
        for(int i=0;i<=2;i++){now=i;build(1,1,n+1);} mp.clear();
        v.push_back(0);
        sort(v.begin(),v.end());v.erase(unique(v.begin(),v.end()),v.end());
        for(int i=0;i<v.size();i++) mp[v[i]]=i+1;
        // dp[0]=-1e9;
        for(int i=0;i<=2;i++){now=i;change(mp[0],1,0);}
        now=0;
        for(int i=1;i<=n;i++){
            int pos=mp[s[i]];
            now=0;dp[i]=max(dp[i],query(1,1,pos-1)+i);
            now=1;dp[i]=max(dp[i],query(1,pos,pos));
            now=2;dp[i]=max(dp[i],query(1,pos+1,n+1)-i);
            now=0;change(pos,1,dp[i]-i);
            now=1;change(pos,1,dp[i]);
            now=2;change(pos,1,dp[i]+i);
        }
        cout<<dp[n]<<endl;
    }
}