topk_sum

[Astral-23]

struct topk_sum {  // min
    size_t k;
    ll sum = 0;
    maxheap<ll> div1;  // min
    minheap<ll> div2;  // min
    // minheap<ll> div1;//MAX
    // maxheap<ll> div2;//MAX

    topk_sum(int k = 0) : k(k) {};

  private:
    void balance() {
        if (k < 0) k = 0;
        // まず個数を合わせる
        while (!div2.empty() && div1.size() < k) {  // サイズが小さい時
            ll in = div2.top();
            div2.pop();
            div1.push(in);
            sum += in;
        }
        while (div1.size() > k) {  // サイズが大きい時
            ll out = div1.top();
            div1.pop();
            sum -= out;
            div2.push(out);
        }

        while (!div1.empty() && !div2.empty()) {
            ll out = div1.top();
            ll in = div2.top();
            if (out <= in) break;  // min
            // if(out >= in) break;//MAX
            div2.pop();
            div1.pop();
            sum += in, sum -= out;
            div1.push(in);
            div2.push(out);
        }
    }

  public:
    void add(ll x) {
        div2.push(x);
        balance();
    }

    void resize(ll new_k) {
        k = new_k;
        balance();
    }

    ll val() { return sum; }

    ll cnt() { return div1.size(); }
    bool full() { return div1.size() == k; }
};

非破壊復元

struct topk_sum {  // max
    size_t k;
    ll sum = 0;
    multiset<ll> div1;
    multiset<ll> div2;

    topk_sum(int k = 0) : k(k) {};

  private:
    void balance() {
        if (k < 0) k = 0;
        // まず個数を合わせる
        while (!div2.empty() && div1.size() < k) {  // サイズが小さい時
            auto in_itr = prev(div2.end());
            div1.insert(*in_itr);
            sum += *in_itr;
            div2.erase(in_itr);
        }
        while (div1.size() > k) {  // サイズが大きい時
            auto out_itr = div1.begin();
            sum -= *out_itr;
            div2.insert(*out_itr);
            div1.erase(out_itr);
        }

        while (!div1.empty() && !div2.empty()) {
            // auto out_itr = div1.begin();     //MAX
            // auto in_itr = prev(div2.end());  //MAX
            // if (*out_itr >= *in_itr) break;  //MAX
            auto out_itr = prev(div1.end());  // min
            auto in_itr = div2.begin();       // min
            if (*out_itr <= *in_itr) break;   // min
            div2.insert(*out_itr);
            div1.insert(*in_itr);
            sum += *in_itr;
            sum -= *out_itr;
            div1.erase(out_itr);
            div2.erase(in_itr);
        }
    }

  public:
    void add(ll x) {
        div2.insert(x);
        balance();
    }

    void del(ll x) {
        if (div1.find(x) != div1.end()) {
            div1.erase(div1.find(x));
            sum -= x;
        } else {
            div2.erase(div2.find(x));
        }
        balance();
    }

    void resize(ll new_k) {
        k = new_k;
        balance();
    }

    ll val() { return sum; }

    vec<ll> dump() {
        vec<ll> res;
        for (ll x : div1) res.push_back(x);
        return res;
    }

    ll cnt() { return div1.size(); }
    bool full() { return div1.size() == k; }
};

https://atcoder.jp/contests/codequeen2024-final-N9tn8QqD/submissions/56471553 https://atcoder.jp/contests/codequeen2024-final-N9tn8QqD/submissions/56471596