#include "partition.h"
#include<algorithm>
#include<set>
#include<iostream>
using namespace std;

set<int> taken;
std::vector<int> solve(std::vector<int> A) {
    taken.clear();
    vector<pair<int, int>> index_paired;
    long long tgsum = 0;
    long long cursum = 0;

    int cnt = A.size()/2 + A.size() % 2;

    for(int i = 0; i < A.size(); i++) {
        index_paired.push_back({A[i], i});
        tgsum += A[i];
    }
    tgsum += tgsum % 2;
    tgsum /= 2;


    sort(index_paired.begin(), index_paired.end());

    //for(auto [n, _] : index_paired) cout << n << " ";
    //cout << endl;
    //for(auto [_, n] : index_paired) cout << n << " ";
    //cout << endl;

    vector<int> subset;
    int rb = A.size()-1;
    int lb = 0;
    int border_ind = -1;
    int offset = index_paired[lb].first;
    for(int i = 0; i < cnt; i++) {
        if(index_paired[rb].first + cursum >= tgsum - offset) {
            cursum += index_paired[lb].first;
            //cout << "Took small number " << index_paired[lb].first << "; cursum = " << cursum << endl;
            if(border_ind < 0) {
                border_ind = i; // first small number in subset
            }
            //taken.insert(lb);
            subset.push_back(lb++);
        } else {
            cursum += index_paired[rb].first;
            //cout << "Took large number " << index_paired[rb].first << "; cursum = " << cursum << endl;
            //taken.insert(rb);
            subset.push_back(rb--);
        }

    }

    // < means we need to shift the largest small number upwards, > means the smallest large number downwards
    if(cursum < tgsum) {
        while(cursum < tgsum) {
            //taken.erase(subset.back());
            //cout << "Shifting small number at i = " << subset.back();
            cursum -= index_paired[subset.back()++].first;
            //taken.insert(subset.back());
            cursum += index_paired[subset.back()].first;
            //cout << "; cursum = " << cursum << endl;
        }
    } else if(border_ind > 0){
        while(cursum - index_paired[subset[border_ind-1]].first + index_paired[subset[border_ind-1]-1].first >= tgsum) {
            //cout << "Shifting large number at i = " << subset[border_ind-1];
            cursum -= index_paired[subset[border_ind-1]--].first;
            cursum += index_paired[subset[border_ind-1]].first;
            //cout << "; cursum = " << cursum << endl;
        }
    }

    vector<int> indices;
    for(int i : subset) {
        indices.push_back(index_paired[i].second);
        taken.insert(index_paired[i].second);
    }
    for(int i = 0; i < A.size(); i++) {
        if(!taken.count(i)) indices.push_back(i);
    }
    return indices;
}
