Sequence Partition

#include <algorithm>
#include <iostream>
#include <map>
#include <set>

#define INF 500000000
#define MAX_N 15100
#define RIGHT 32768
#define SIZE 65537

using namespace std;

int N, K, T;
int A[MAX_N], S[MAX_N];
map<int,int> Idx; // Contains index in Tree of prefix sum

// Tree class
// first contains min partitions so that each partition sum is <= M
// second contains max partitions that can be made so that each partition sum is <= M

class T
{
private:
    pair<int,int> node[SIZE];
public:
    pair<int,int> query(int r, int n, int a, int b);
    void reset();
    void update(int i, pair<int,int> p, int n, int a, int b);
} Tree;

void T::reset()
{
    for(int i = 0; i < SIZE; i++)
        node[i].first = INF, node[i].second = -INF;

    update(Idx[0], make_pair(0,0), 1, 1, RIGHT);
}

void T::update(int i, pair<int,int> p, int n, int a, int b)
{
    if(i < a || i > b)
        return;
    if(i == a && i == b)
    {
        node[n] = p;
        return;
    }

    int mid = (a+b) / 2;
    update(i, p, 2*n, a, mid);
    update(i, p, 2*n + 1, mid + 1, b);
    node[n].first = min(node[2*n].first, node[2*n + 1].first);
    node[n].second = max(node[2*n].second, node[2*n + 1].second);
}

pair<int,int> T::query(int r, int n, int a, int b)
{
    if(b < r)
        return make_pair(INF, -INF);
    if(a >= r)
        return node[n];

    int mid = (a+b) / 2;
    pair<int,int> p = query(r, 2*n, a, mid);
    pair<int,int> q = query(r, 2*n + 1, mid + 1, b);
    return make_pair(min(p.first, q.first), max(p.second, q.second));
}

bool possible(int m)
{
    // Compress prefix sums

    set<int> s;
    for(int i = 0; i <= N; i++)
        s.insert(S[i]), s.insert(S[i] - m);

    Idx.clear();
    int i = 1;
    for(set<int>::iterator it = s.begin(); it != s.end(); it++, i++)
        Idx[*it] = i;

    // Find answer with DP

    Tree.reset();

    for(int i = 1; i <= N; i++)
    {
        pair<int,int> p = Tree.query(Idx[S[i] - m], 1, 1, RIGHT);
        p.first++, p.second++;
        Tree.update(Idx[S[i]], p, 1, 1, RIGHT);
    }

    pair<int,int> p = Tree.query(Idx[S[N]], 1, 1, RIGHT);
    return p.first <= K && p.second >= K;
}

int minM()
{
    // Binary search on M

    int a = -INF, b = INF;
    int p = (a+b) / 2;

    while(b > a+1)
    {
        if(possible(p))
            b = p;
        else
            a = p;

        p = (a+b) / 2;
    }

    if(possible(a))
        return a;
    else
        return b;
}

int main()
{
    cin >> T;

    while(T--)
    {
        cin >> N >> K;
        int s = 0;

        for(int i = 1; i <= N; i++)
        {
            cin >> A[i];
            S[i] = S[i-1] + A[i];
        }

        cout << minM() << "\n";
    }
}