IOI '17 P4 - The Big Prize

#include <iostream>
#include <vector>
#include <algorithm>
#include <set>
#include <cassert>
#include "prize.h"

using namespace std;

vector<int> l, r;
set<int> q; // queried indices that were not expensive
int expensiveCnt = 0;
void query(int i) {
    if (l[i] == -1) {
        vector<int> p = ask(i);
        l[i] = p[0], r[i] = p[1];
        if (l[i] + r[i] == expensiveCnt) q.insert(i); // if not expensive
    }
}

bool isExpensive(int i) {
    if (l[i] == -1) query(i);
    return l[i] + r[i] < expensiveCnt;
}
bool isDiamond(int i) {
    return l[i] + r[i] == 0;
}

int diamond = -1;
void divideAndConquer(int left, int right, int k, int lBefore = 0, int rAfter = 0) { // k is how many expensives in this range
    if (!k || diamond != -1) return;

    set<int>::iterator it = q.lower_bound(left);
    if (it != q.end() && *it <= right) {
        int mid = *it;
        divideAndConquer(left, mid-1, l[mid]-lBefore, lBefore, r[mid]);
        divideAndConquer(mid+1, right, r[mid]-rAfter, l[mid], rAfter);
        return;
    }

    int mid = (left + right) / 2, startMid = mid;
    bool reachedLeft = false;
    while (isExpensive(mid)) {
        if (isDiamond(mid)) {
            diamond = mid;
            return;
        }
        if (!reachedLeft) {
            --mid;
            if (mid < left) reachedLeft = true, mid = startMid;
        }
        if (reachedLeft) {
            ++mid;
            if (mid > right) {
                assert(k == right-left+1);
                return;
            }
        }
    }

    assert(!isExpensive(mid));
    if (!reachedLeft) divideAndConquer(left, mid-1, l[mid]-lBefore, lBefore, r[mid]);
    divideAndConquer(mid+1, right, r[mid]-rAfter, l[mid], rAfter);
}

int find_best(int n){
    l = r = vector<int>(n, -1);
    // we will call all objects that are not the cheapest objects expensive
    int maxExpensive = min(3, n), startQ = min(maxExpensive+1, n);
    for (int i = 0; i < startQ; ++i) {
        query(i);
        if (isDiamond(i)) return i;
        expensiveCnt = max(expensiveCnt, l[i]+r[i]);
    }

    divideAndConquer(0, n-1, expensiveCnt);
    return diamond;
}