Codeforces Round #783 (Div. 2) D - Optimal Partition(dp/权值线段树 2100)

题意：

问我们将数组$a$分割成任意非空连续段，并且每一段的价值满足：

• 若所有元素的和为正数，则价值等于子数组的长度

• 若所有元素的和为 $0$ ，则价值等于 $0$

• 若所有元素的和为正数，则价值等于子数组的长度的相反数。

的情况下总价值最大是多少。

• 容易想到$n^2$的朴素$dp$做法：

• 一维循环$i$

• 二维更新每个$dp[i]$使其最大，可以从$dp[j]+val((j+1)$~$i)$取最大转移。

void bf(){
    for(int i=1;i<=n;i++)pre[i]=pre[i-1]+a[i];
    for(int i=1;i<=n;i++){
        dp[i]=-1e9;
        for (int j=0;j<i;j++){//更新转移取MAX
            if (pre[i]-pre[j]>0)f[i]=max(dp[i],dp[j]+(i-j));
            if (pre[i]==pre[j])f[i]=max(dp[i],dp[j]);
            if (pre[i]-pre[j]<0)f[i]=max(dp[i],dp[j]+(j-i));
        }
    }
    cout<<dp[n]<<endl;
}

• 但是由于$n$的上限是$5 * 10^5$所以$n^2$必定会$tle$。因此我们需要进行优化。

• 通过公式的推导我们可以发现：
  $1:dp[j]+(i-j)=dp[j]-j+i$
  $2:dp[j]=dp[j]$
  $3:dp[j]+(j-i)=dp[j]+j-i$
  由于$i$ 是不变的，所以我们每次更新最大的$dp[j]$显然只用求上面三式在满足各自限制的前提下最大的值即可：
  $pre[j]$<$pre[i]：$ 进行$1$式更新。
  $pre[j]$==$pre[i]：$ 进行$2$式更新。
  $pre[j]$>$pre[i]：$ 进行$3$式更新。
  即我们可以用权值线段树进行维护。也就是用线段树的下标来表示$pre[j]$，那么查询操作同样可以理解为：
  $query(1,pre[j]-1)$
  $query(pre[j],pre[j])$
  $query(pre[j]+1,n)$

参考代码+注释:

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N = 5e5+10;
const int INF = 1e9;
int a[N];//前缀和 
int dp[N];
vector<int>vec;
struct Node{
    int l,r;
    int ma1,ma2,ma3;//表示三种转移状态 
}tr[N*4];
void pushup(Node& root,Node& left,Node& right){
    root.ma1=max(left.ma1,right.ma1);
    root.ma2=max(left.ma2,right.ma2);
    root.ma3=max(left.ma3,right.ma3);
}
void pushup(int u){
    pushup(tr[u],tr[u<<1],tr[u<<1|1]);
}
int find(int x){
    return lower_bound(vec.begin(),vec.end(),x)-vec.begin()+1;//下标从1开始 
}
void build(int u,int l,int r){
    tr[u]={l,r,-INF,-INF,-INF};
    if(l==r)return ;
    int mid=l+r>>1;
    build(u<<1,l,mid),build(u<<1|1,mid+1,r);
}
void modify(int u,int x,int v,int d){
    if(tr[u].l==x&&tr[u].r==x){
        tr[u].ma1=max(tr[u].ma1,v-d);
        tr[u].ma2=max(tr[u].ma2,v);
        tr[u].ma3=max(tr[u].ma3,v+d);
    }
    else{
        int mid=tr[u].l+tr[u].r>>1;
        if(x<=mid)modify(u<<1,x,v,d);
        else if(x>mid)modify(u<<1|1,x,v,d);
        pushup(u);
    }
}
Node query(int u,int l,int r){
    if(tr[u].l>=l&&tr[u].r<=r){
        return tr[u];
    }
    int mid=tr[u].l+tr[u].r>>1;
    if(r<=mid)return query(u<<1,l,r);
    else if(l>mid)return query(u<<1|1,l,r);
    else{
        Node uu,left,right;
        left=query(u<<1,l,mid);
        right=query(u<<1|1,mid+1,r);
        pushup(uu,left,right);
        return uu;
    }
}
void solve(){
    int n;
    cin>>n;
    vec.clear();//清空
    a[0]=0;//因为权值线段树所以每次测试不一定为0 
    for(int i=1;i<=n;i++)cin>>a[i],a[i]+=a[i-1],vec.push_back(a[i]);
    vec.push_back(0);
    sort(vec.begin(),vec.end());
    vec.erase(unique(vec.begin(),vec.end()),vec.end());//离散化
    for(int i=0;i<=n;i++){
        a[i]=find(a[i]);
    }
    int m=vec.size();
    build(1,1,m);//建树
    dp[0]=0;//dp入口
    modify(1,a[0],0,0);//dp入口
    for(int i=1;i<=n;i++){
        dp[i]=-INF;
        if(a[i]>1){
            Node u=query(1,1,a[i]-1);//1式
            dp[i]=max(dp[i],u.ma1+i);
        }
        if(a[i]<m){
            Node u=query(1,a[i]+1,m);//3式
            dp[i]=max(dp[i],u.ma3-i);
        }
        Node u=query(1,a[i],a[i]);
        dp[i]=max(dp[i],u.ma2);//2式
        modify(1,a[i],dp[i],i);//修改供后面更新
    }
    cout<<dp[n]<<endl; 
}
signed main(){
    int t;
    cin>>t;
    while(t--){
        solve();
    }
    return 0;
}