CUBE-T197: Fair and Square

#include <bits/stdc++.h>
using namespace std;

typedef long long lli;

const int N = 1500;
const int K = 50;

int arr[N + 1];
lli cost[N + 1][N + 1], dp[N + 1][K + 1], sumL, sumR;
priority_queue<int> l;
priority_queue<int, vector<int>, greater<int>> r;

void medSetup(){
    while(!l.empty()){
        l.pop();
    }
    while(!r.empty()){
        r.pop();
    }
    sumL = 0;
    sumR = 0;
}

void medInsert(int x){
    if(l.empty() || x < l.top()){
        if(((int)l.size() + r.size()) % 2 == 1){
            sumR += l.top();
            r.push(l.top());
            sumL -= l.top();
            l.pop();
        }
        sumL += x;
        l.push(x);
    } else if(r.empty() || x > r.top()){
        if(((int)l.size() + r.size()) % 2 == 0){
            sumL += r.top();
            l.push(r.top());
            sumR -= r.top();
            r.pop();
        }
        sumR += x;
        r.push(x);
    } else {
        if(((int)l.size() + r.size()) % 2 == 1){
            sumR += x;
            r.push(x);
        } else {
            sumL += x;
            l.push(x);
        }
    }
}
lli medMinCost(){
    if(((int)l.size() + r.size()) % 2 == 1){
        return (lli)l.size() * l.top() - sumL + sumR - (lli)r.size() * l.top();
    } else {
        lli ansL = (lli)l.size() * l.top() - sumL + sumR - (lli)r.size() * l.top();
        lli ansR = (lli)l.size() * r.top() - sumL + sumR - (lli)r.size() * r.top();
        return min(ansL, ansR);
    }
}

int main(){

    int n, parts;
    scanf("%d%d", &n, &parts);
    for(int i = 1; i <= n; ++i){
        scanf("%d", &arr[i]);
    }
    for(int i = 1; i <= n; ++i){
        medSetup();
        for(int j = i; j <= n; ++j){
            medInsert(arr[j]);
            cost[i][j] = medMinCost();
        }
    }
    dp[0][0] = 0;
    for(int i = 1; i <= n; ++i){
        dp[i][0] = 1e18;
    }
    for(int p = 1; p <= parts; ++p){
        dp[0][p] = 1e18;
    }
    for(int i = 1; i <= n; ++i){
        for(int p = 1; p <= parts; ++p){
            lli mn = 1e18;
            for(int j = 1; j <= i; ++j){
                mn = min(mn, cost[j][i] + dp[j - 1][p - 1]);
            }
            dp[i][p] = mn;
        }
    }
    cout << dp[n][parts];

    return 0;
}