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

std::vector<int> solve(std::vector<int> A) {
    long long sum = 0;
    for (const auto x : A) sum += x;

    std::vector<int> p(A.size());
    for (int i = 0; i < A.size(); i++) p[i] = i;
    std::sort(p.begin(), p.end(), [&A] (const auto a, const auto b) { return A[a] < A[b]; });

    struct solution {
        int l, r;
        int diff;
    };

    solution best = { .diff = INT_MAX };
    int l = 0, r = 0;
    long long s = 0;
    while (true) {
        if (s >= sum / 2 + sum % 2) {
            const auto diff = abs((r - l) * 2 - (int)A.size());
            if (diff < best.diff) best = { .l = l, .r = r, .diff = diff };
        }

        if (s < sum / 2 + sum % 2 && r < A.size()) s += A[p[r++]];
        else if (l < r) s -= A[p[l++]];
        else break;
    }

    std::set<int> p1, p2;
    for (int i = 0; i < best.l; i++) p2.insert(p[i]);
    for (int i = best.l; i < best.r; i++) p1.insert(p[i]);
    for (int i = best.r; i < A.size(); i++) p2.insert(p[i]);

    std::vector<int> ans;
    ans.insert(ans.end(), p1.begin(), p1.end());
    ans.insert(ans.end(), p2.begin(), p2.end());
    return ans;
}
