CUBE-197: Fair and Square

#include <bits/stdc++.h>

using namespace std;
const int K = 50 + 10;
const int N = 1500 + 10;
const int INF = 1e9;
using pi = pair <int, int>;
int tree[3][N];
int sz, sum;
int ar[N], med[N][N], dp[K][N];
vector <int> coor;

int Index(int x){
    int l = 1, r = sz - 1;
    while(l <= r){
        int mid = (l + r) / 2;
        if(coor[mid] == x) return mid;
        else if(coor[mid] < x) l = mid + 1;
        else r = mid - 1;
    }
}

void Update(int pst, int val1, int val2){
    for(int i = pst; i < sz; i += i & -i){
        tree[1][i] += val1;
        tree[2][i] += val2;
    }
}

pi Query(int pst, pi q = {0, 0}){
    for(int i = pst; i > 0; i -= i & -i){
        q.first += tree[1][i];
        q.second += tree[2][i];
    }
    return q;
}

void Median(int i, int j, int len){
    int mp = (len / 2) + (len % 2), ms = INF, mv = 0, mc = INF, mn = INF;
    int l = 1, r = sz - 1;
    while(l <= r){
        int mid = (l + r) / 2;
        pi q = Query(mid);
        int bsc = q.first, bss = q.second;
        if(bsc >= mp){
            r = mid - 1;
            ms = min(ms, bss);
            mc = min(mc, bsc);
            mn = min(mn, mid);
        }
        else l = mid + 1;
    }
    mv = coor[mn];
    med[i][j] = ( mc * mv - ms ) + ( (sum - ms) - (len - mc) * mv );
}

int main(){

    int n, part;
    scanf("%d%d", &n, &part);

    for(int i=1;i<=n;i++){
        scanf("%d", &ar[i]);
        coor.push_back(ar[i]);
    }

    sort(coor.begin(), coor.end());
    coor.resize( unique(coor.begin(), coor.end()) - coor.begin() );
    coor.insert(coor.begin(), 0);
    sz = coor.size(); // for(int i=1;i<sz;i++) printf("%d ", coor[i]);

    for(int l=1;l<=n;l++){
        int i, j; /// [i,j] l = j-i+1
        for(j=l;j<=n;j++){ i = j - l + 1;
            if(i > 1) Update(Index(ar[i-1]), -1, -ar[i-1]), sum -= ar[i-1];
            Update(Index(ar[j]), 1, ar[j]), sum += ar[j];
            Median(i, j, l);
        }
        l ++;
        for(i=n-l+1;i>=1;i--){ j = i + l - 1;
            if(j < n) Update(Index(ar[j+1]), -1, -ar[j+1]), sum -= ar[j+1];
            Update(Index(ar[i]), 1, ar[i]), sum += ar[i];
            Median(i, j, l);
        }
    }

    dp[0][0] = 0;
    for(int i=1;i<=n;i++) dp[0][i] = INF;

    for(int k=1;k<=part;k++){
        for(int j=1;j<=n;j++){
            dp[k][j] = INF;
            for(int i=1;i<=j;i++){
                dp[k][j] = min(dp[k][j], dp[k-1][i-1] + med[i][j]);
            }
        }
    }

    printf("%d", dp[part][n]);

    return 0;
}