cpp_src/other/TopKset.hpp

View this file on GitHub
Last update: 2023-12-31 13:48:59+09:00
Include: #include "cpp_src/other/TopKset.hpp"
Link: https://atcoder.jp/contests/abc281/submissions/37651034

Code

// ABC281E : https://atcoder.jp/contests/abc281/submissions/37651034
// sl ... to kth smallest elements

void maintain_topk() {
    multiset<ll> sl, sr;
    ll suml = 0, sumr = 0;

    auto del = [&](ll x) {
        if (sl.find(x) != sl.end()) {
            sl.erase(sl.find(x));
            suml -= x;
        } else {
            sr.erase(sr.find(x));
            sumr -= x;
        }
    };

    auto add = [&](ll x) {
        if (!(SZ(sr) != 0 || *sr.begin() <= x)) {
            sl.insert(x);
            suml += x;
        } else {
            sr.insert(x);
            sumr += x;
        }
    };

    auto balance = [&]() {
        while (SZ(sl) < K && SZ(sr) > 0) {
            ll f = *sr.begin();
            sr.erase(sr.begin());
            sl.insert(f);
            suml += f;
            sumr -= f;
        }
        while (SZ(sl) > K) {
            ll f = *(--sl.end());
            sl.erase(--sl.end());
            sr.insert(f);
            suml -= f;
            sumr += f;
        }
    };
}

#line 1 "cpp_src/other/TopKset.hpp"
// ABC281E : https://atcoder.jp/contests/abc281/submissions/37651034
// sl ... to kth smallest elements

void maintain_topk() {
    multiset<ll> sl, sr;
    ll suml = 0, sumr = 0;

    auto del = [&](ll x) {
        if (sl.find(x) != sl.end()) {
            sl.erase(sl.find(x));
            suml -= x;
        } else {
            sr.erase(sr.find(x));
            sumr -= x;
        }
    };

    auto add = [&](ll x) {
        if (!(SZ(sr) != 0 || *sr.begin() <= x)) {
            sl.insert(x);
            suml += x;
        } else {
            sr.insert(x);
            sumr += x;
        }
    };

    auto balance = [&]() {
        while (SZ(sl) < K && SZ(sr) > 0) {
            ll f = *sr.begin();
            sr.erase(sr.begin());
            sl.insert(f);
            suml += f;
            sumr -= f;
        }
        while (SZ(sl) > K) {
            ll f = *(--sl.end());
            sl.erase(--sl.end());
            sr.insert(f);
            suml -= f;
            sumr += f;
        }
    };
}