Submission

Score: 76

User: qweqwe

Problemset: แคง (Kang)

Language: cpp

Time: 2.087 second

Submitted On: 2026-03-26 22:06:30

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using pii = pair<ll,ll>;

struct Node {
    bool inBig;
    set<pii>::iterator it;
};

set<pii> small, big;
unordered_map<ll, int> cnt;
unordered_map<ll, Node> pos;
ll sum_big = 0, sum = 0;

inline void erase_val(ll v) {
    auto itp = pos.find(v);
    if (itp == pos.end()) return;

    if (itp->second.inBig) {
        sum_big -= itp->second.it->first;
        big.erase(itp->second.it);
    } else {
        small.erase(itp->second.it);
    }
    pos.erase(itp);
}

inline void insert_small(ll v) {
    auto it = small.insert({1LL * v * cnt[v], v}).first;
    pos[v] = {false, it};
}

inline void add(ll v) {
    if (cnt[v]) erase_val(v);
    ++cnt[v];
    sum += v;
    insert_small(v);
}

inline void move_small_to_big() {
    auto it = prev(small.end());
    pii x = *it;
    small.erase(it);
    auto jt = big.insert(x).first;
    pos[x.second] = {true, jt};
    sum_big += x.first;
}

inline void move_big_to_small() {
    auto it = big.begin();
    pii x = *it;
    sum_big -= x.first;
    big.erase(it);
    auto jt = small.insert(x).first;
    pos[x.second] = {false, jt};
}

inline void balance(int k) {
    while ((int)big.size() < k && !small.empty()) {
        move_small_to_big();
    }
    while ((int)big.size() > k) {
        move_big_to_small();
    }

    while (!small.empty() && !big.empty()) {
        auto its = prev(small.end());
        auto itb = big.begin();

        if (its->first > itb->first) {
            pii s = *its, b = *itb;
            sum_big += s.first - b.first;

            small.erase(its);
            big.erase(itb);

            auto jt1 = small.insert(b).first;
            auto jt2 = big.insert(s).first;
            pos[b.second] = {false, jt1};
            pos[s.second] = {true, jt2};
        } else {
            break;
        }
    }
}

vector<long long> capsize(vector<int> a, vector<int> b) {
    cnt.reserve(a.size() + b.size() + 10);
    pos.reserve(a.size() + b.size() + 10);

    vector<ll> ans;
    ans.reserve(b.size());

    for (int x : a) add(x);

    for (int i = 0; i < (int)b.size(); i++) {
        add(b[i]);
        balance(i + 1);
        ans.push_back(sum - sum_big);
    }
    return ans;
}
/*
5 4
6 3 9 6 3
4 6 9 6

19 10 4 0

8 7
1 2 3 4 5 6 7 8
9 10 11 12 13 14 15

36 36 36 36 36 36 36

5 4
32 78 4 10 9
52 46 33 77

107 101 88 88

5 4
1 2 3 4 5
4 4 4 4
*/