#include "partition.h"
#include <algorithm>
#include <iostream>

using namespace std;

vector<int> solve(vector<int> a) {
    int n = a.size();
    if(n == 1) {
        return {0};
    }

    vector<int> idx(n);
    for(int i = 0; i < n; i++) {
        idx[i] = i;
    }
    sort(idx.begin(), idx.end(), [&](int x, int y) {
        return a[x] > a[y];
    });
    long long sum = 0;
    for(int i = 0; i < n; i++) {
        sum += a[i];
    }

    long long snow = sum;
    int ptr = 1, best = n + 1, besti = 0, bestj = 0;
    for(int i = 0; i < n; i++) {
        if(i > 0) {
            snow += a[idx[i]];
        }
        while(ptr < n && snow - a[idx[ptr]] > sum / 2) {
            snow -= a[idx[ptr++]];
        }
        int cnt = i + 1 + n - ptr;
        if(abs(2 * cnt - n) < best) {
            best = abs(2 * cnt - n);
            besti = i;
            bestj = ptr;
        }
    }

    vector<int> answer;
    answer.reserve(n);
    for(int i = 0; i <= besti; i++) {
        answer.push_back(idx[i]);
    }
    for(int i = bestj; i < n; i++) {
        answer.push_back(idx[i]);
    }
    for(int i = besti + 1; i < bestj; i++) {
        answer.push_back(idx[i]);
    }
    return answer;
}
