配列をソートしたときに先頭からK番目までの和を求めるやつ

　ときどき必要になるけど、実装がつらいですよね。

　AWC0041-D で必要になったので、この機会にライブラリを整備しておきます。まあ、AI ポン出しなのですが。

struct BestKSum {
    int k;
    multiset<long long> in, out;
    long long best = 0;
    long long other = 0;

    BestKSum(int k_) : k(k_) {}

    void fix() {
        while ((int)in.size() > k) {
            auto it = in.begin();
            long long x = *it;
            in.erase(it);
            best -= x;
            out.insert(x);
            other += x;
        }
        while ((int)in.size() < k && !out.empty()) {
            auto it = prev(out.end());
            long long x = *it;
            out.erase(it);
            other -= x;
            in.insert(x);
            best += x;
        }
        while (!in.empty() && !out.empty()) {
            auto a = in.begin();
            auto b = prev(out.end());
            if (*a >= *b) break;
            long long x = *a, y = *b;
            in.erase(a);
            out.erase(b);
            in.insert(y);
            out.insert(x);
            best += y - x;
            other += x - y;
        }
    }

    void insert(long long x) {
        if (!in.empty() && x >= *in.begin()) {
            in.insert(x);
            best += x;
        } else {
            out.insert(x);
            other += x;
        }
        fix();
    }

    void erase(long long x) {
        auto it = in.find(x);
        if (it != in.end()) {
            in.erase(it);
            best -= x;
        } else {
            it = out.find(x);
            assert(it != out.end());
            out.erase(it);
            other -= x;
        }
        fix();
    }

    void change(long long old_x, long long new_x) {
        erase(old_x);
        insert(new_x);
    }

    long long sum_best() const { return best; }
    long long sum_other() const { return other; }
    long long sum_all() const { return best + other; }
};

使い方

• BestKSum(int k_) : k(k_) {} コンストラクタ。上位何個を持つかを決めます
• insert(x) 値 x を $1$ 個追加します
• erase(x) 値 x を $1$ 個削除します
• change(old_x, new_x) 値を $1$ 個更新します
• sum_best() 現在の上位 $K$ 個の和を返します。要素数が $k$ 未満なら、存在する全要素の和になります
• sum_other() 現在の上位 $K$ 個以外の和を返します

使用例

ABC306-E

int main(){
    int n,k,q;
    cin >> n >> k >> q;
    vector<int> a(n);
    BestKSum bks(k);
    rep(i,n){
        bks.insert(0);
    }
    while(q--){
        int x,y;
        cin >> x >> y;
        x--;
        bks.change(a[x],y);
        a[x] = y;
        cout << bks.sum_best() << endl;
    }
}