CF783-D Optimal Partition

Optimal Partition

// dp + fenwick + 离散 + 前缀和
// 可以证明:对于最优划分，分数为0和负数的区间一定可以优化成长度只有1的子段
// 现在只考虑分数为正数的情况：f[i] = max{f[j] - j} + i; 则维护一棵Fenwick Tree即可
// f[i] = max(f[i - 1] + (a[i] < 0 ? -1 : 0), fen.sum(s[i]));
#include <bits/stdc++.h>

using namespace std;
using ll = long long;

const int INF = 1e9;

struct Fenwick {
    const int n;
    std::vector<int> a;
    Fenwick(int n) : n(n), a(n, -INF) {}
    void add(int x, int v) {
        for (int i = x + 1; i <= n; i += i & -i) {
            a[i - 1] = std::max(a[i - 1], v);
        }
    }
    int sum(int x) {
        int ans = -INF;
        for (int i = x; i > 0; i -= i & -i) {
            ans = std::max(ans, a[i - 1]);
        }
        return ans;
    }
};

void solve(){
    int n; cin >> n;
    
    vector<int>a(n);
    for(int i = 0; i < n; ++i) cin >> a[i];
    
    vector<ll>s(n + 1);
    for(int i = 1; i <= n; ++i) s[i] = s[i - 1] + a[i - 1];
    
    auto v = s;
    sort(v.begin(), v.end());
    for(int i = 0; i <= n; ++i) s[i] = lower_bound(v.begin(), v.end(), s[i]) - v.begin();
    
    Fenwick fen(n + 1);
    vector<int>f(n + 1);
    for(int i = 1; i <= n; ++i) {
        fen.add(s[i - 1], f[i - 1] - i + 1);
        f[i] = max(f[i - 1] + (a[i - 1] < 0 ? -1 : 0), i + fen.sum(s[i]));
    }
    cout << f[n] << endl;
}

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