#include "partition.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 i, l, r;
        long long diff;

        bool operator>(const solution &other) const { return diff < other.diff; }
        bool operator<(const solution &other) const { return diff > other.diff; }
    };

    long long pref = 0;
    solution best = { .diff = LLONG_MAX };
    for (int i = 0; i <= A.size(); i++) {
        if (i > 0) pref += A[p[i - 1]];
        int l = i, r = l;
        long long s = pref;
        const auto diff = [&s, sum] () { return (s * 2 - sum); };
        while (l <= r && r <= A.size()) {
            if (diff() >= 0 && diff() < best.diff) {
                best = {
                    .i = i, .l = l, .r = r,
                    .diff = diff()
                };
//                std::cout << best.diff << ' ' << best.i << ' ' << best.l << ' ' << best.r << ' ' << s << ' ' << sum << '\n';
            }

            if (diff() < 0 && 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.i; i++) p1.insert(p[i]);
    for (int i = best.i; 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;
}

/*
1 0 3 8 2 6 7 4 5
6 5 3 3 6 4 4 1 1

5 6 6 3 1 1 4 4 3
1 1 3 3 4 4 5 6 6
6 5 4 1 1 | 3 3 4 6

2 7 3 6 4 0 1 5
3 3 2 2 3 2 2 2

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