#include "partition.h"
#include <iostream>
#include <algorithm>
#include <queue>
#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]; });

    std::set<std::pair<int, int>> bag;
    std::priority_queue<std::pair<int, int>> rem;
    long long s = 0;
    for (const auto i : p) {
        if (s + A[i] >= sum / 2 + sum % 2) {
            rem.push({A[i], i});
            continue;
        }
        s += A[i];
        bag.insert({A[i], i});
    }

    bag.insert(rem.top());
    s += rem.top().first;
    rem.pop();

    const auto goal = A.size() / 2 + A.size() % 2;
    while (bag.size() > goal) {
        const auto x = *bag.begin();
        s -= x.first;
        bag.erase(bag.begin());

        if (s < sum / 2 + sum % 2) {
            if (rem.empty() || rem.top().first == x.first) {
                bag.insert(x);
                break;
            }
            bag.insert(rem.top());
            s += rem.top().first;
            rem.pop();
        }
        rem.push(x);
    }

    std::vector<int> ans;
    ans.reserve(A.size());
    for (int i = 0; i < A.size(); i++)
        if (bag.count({A[i], i})) ans.push_back(i);
    for (int i = 0; i < A.size(); i++)
        if (!bag.count({A[i], i})) ans.push_back(i);
    return ans;
}

/*
2 2 2 2 2 3 3 3
2 2 3 3
*/
