Donutsプロコンチャレンジ2015 D - ドーナツの箱詰め

問題ページ

まず のときの解き方を考える。 についてソートしておく。この数列の差分について小さい方から 個の和が答えになる。差分が大きいところから 個と一番小さい箱にドーナツを入れることでこれは達成できる。

数列の小さい方から 個の和を求めるのにはセグ木、平衡二分探索木などのデータ構造を用いればよい。今回はRBSTを使う。

番目の箱が潰されるときのデータ構造の変化について考える。 番目の箱より小さい最大の箱を 、 番目の箱より大きい最小の箱を と表すことにする。、 をRBSTから消去し、 をRBSTに追加すればよい。prev, next については std::set で実現可能なのでこの問題は解けた。

▶ソースコード(クリックで展開)

#include <bits/stdc++.h>  
using namespace std;  
using ll = long long;  
using PII = pair<ll, ll>;  
#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)  
#define REP(i, n) FOR(i, 0, n)  
#define ALL(x) x.begin(), x.end()  
template<typename T> void chmin(T &a, const T &b) { a = min(a, b); }  
template<typename T> void chmax(T &a, const T &b) { a = max(a, b); }  
struct FastIO {FastIO() { cin.tie(0); ios::sync_with_stdio(0); }}fastiofastio;  
#ifdef DEBUG  
#include "../program_contest_library/memo/dump.hpp"  
#else  
#define dump(...)  
#endif  
const ll INF = 1LL<<60;  
  
template<class M>  
struct RBST {  
    using T = typename M::T;  
    using E = typename M::E;  
  
    struct node {  
        node *l, *r, *p;  
        T val;  
        E lazy;  
        bool rev;  
        int sz;  
  
        node(T v, E p) : l(nullptr),r(nullptr),val(v),lazy(p),rev(false),sz(1) {}  
    };  
  
    int size(node* a) { return !a ? 0 : a->sz; }  
    T val(node* a) { return !a ? M::dt() : (eval(a), a->val); }  
    node* fix(node* a) {  
        a->sz = size(a->l) + 1 + size(a->r);  
        a->val = M::f(val(a->l), a->val, val(a->r));  
        return a;  
    }  
    void eval(node* a) {  
        if(a->lazy != M::de()) {  
            a->val = M::g(a->val, a->lazy, size(a));  
            if(a->l) a->l->lazy = M::h(a->l->lazy, a->lazy);  
            if(a->r) a->r->lazy = M::h(a->r->lazy, a->lazy);  
            a->lazy = M::de();  
        }  
        if(a->rev) {  
            swap(a->l, a->r);  
            if(a->l) a->l->rev ^= 1;  
            if(a->r) a->r->rev ^= 1;  
            a->rev = false;  
        }  
    }  
  
    inline int xor128() {  
        static int x = 123456789, y = 362436069, z = 521288629, w = 88675123;  
        int t;  
        t = x ^ (x << 11); x = y; y = z; z = w;  
        return w = (w ^ (w >> 19)) ^ (t ^ (t >> 8));  
    }  
  
    node* merge(node *a, node *b) {  
        if(!a) return b;  
        if(!b) return a;  
        eval(a); eval(b);  
        if(xor128() % (size(a) + size(b)) < size(a)) {  
            a->r = merge(a->r, b);  
            if(a->r) a->r->p = a;  
            return fix(a);  
        } else {  
            b->l = merge(a, b->l);  
            if(b->l) b->l->p = b;  
            return fix(b);  
        }  
    }  
    // [0,k) [k,n)  
    pair<node*, node*> split(node *a, int k) {  
        if(!a) return pair<node*, node*>(nullptr, nullptr);  
        eval(a);  
        node *sl, *sr;  
        if(k <= size(a->l)) {  
            tie(sl, sr) = split(a->l, k);  
            if(a->l) a->l->p = nullptr;  
            a->l = sr;  
            if(a->l) a->l->p = a;  
            return pair<node*, node*>(sl, fix(a));  
        }  
        tie(sl, sr) = split(a->r, k - size(a->l) - 1);  
        if(a->r) a->r->p = nullptr;  
        a->r = sl;  
        if(a->r) a->r->p = a;  
        return pair<node*, node*>(fix(a), sr);  
    }  
    // 要素の挿入/削除  
    void insert(node*& a, int k, const T& x) {  
        node *sl, *sr;  
        tie(sl, sr) = split(a, k);  
        a = merge(sl, merge(new node(x, M::de()), sr));  
    }  
    T erase(node*& a, int k) {  
        node *sl, *sr, *tl, *tr;  
        tie(sl, sr) = split(a, k + 1);  
        tie(tl, tr) = split(sl, k);  
        a = merge(tl, sr);  
        return val(tr);  
    }  
    // 区間クエリ  
    T query(node*& a, int l, int r) {  
        node *sl, *sr, *tl, *tr;  
        tie(sl, sr) = split(a, r);  
        tie(tl, tr) = split(sl, l);  
        T res = !tr ? M::dt() : tr->val;  
        a = merge(merge(tl, tr), sr);  
        return res;  
    }  
    // x以上の最小の位置  
    int lower_bound(node *t, const T &x) {  
        if(!t) return 0;  
        if(x <= val(t)) return lower_bound(t->l, x);  
        return lower_bound(t->r, x) + size(t->l) + 1;  
    }  
    // xより大きい最小の位置  
    int upper_bound(node *t, const T &x) {  
        if(!t) return 0;  
        if(x < val(t)) return upper_bound(t->l, x);  
        return upper_bound(t->r, x) + size(t->l) + 1;  
    }  
    // xの個数  
    int count(node *t, const T &x) {  
        return upper_bound(t, x) - lower_bound(t, x);  
    }  
    // 要素xを追加する  
    void insert_key(node *&t, const T &x) {  
        insert(t, lower_bound(t, x), x);  
    }  
    // 要素xを消す  
    void erase_key(node *&t, const T &x) {  
        // if(!count(t, x)) return;  
        erase(t, lower_bound(t, x));  
    }  
    // デバッグ用  
    void debug(node* t) {  
        if(t == nullptr) return;  
        cerr << "{";  
        debug(t->l);  
        cerr << " " << t->val << " ";  
        debug(t->r);  
        cerr << "}";  
    }  
};  
  
struct update_sum {  
    using T = PII;  // (頂点の値, 部分木の総和)  
    using E = ll;  
    static T dt() { return PII(0, 0); }  
    static constexpr E de() { return 0; }  
    static T f(const T &l, const T &self, const T &r) {  
        return PII(self.first, l.second + self.first + r.second);  
    }  
    static T g(const T &a, const E &b, const int &sz) {  
        return b == de() ? a : PII(b, b);  
    }  
    static E h(const E &a, const E &b) {  
        return b == de() ? a : b;  
    }  
};  
  
int main(void) {  
    ll n, k;  
    cin >> n >> k;  
    vector<ll> a(n);  
    REP(i, n) cin >> a[i];  
  
    set<PII> st;  
    REP(i, n) st.insert({a[i], i});  
  
    // 小さい方からx個の合計を求めたい  
    vector<ll> sa(a);  
    sort(ALL(sa));  
    RBST<update_sum> tree;  
    RBST<update_sum>::node *root = nullptr;  
    FOR(i, 1, n) {  
        ll diff = sa[i]-sa[i-1];  
        tree.insert_key(root, PII(diff, diff));  
    }  
  
    cout << tree.query(root, 0, n-k).second << endl;  
    ll q;  
    cin >> q;  
    REP(i, q) {  
        ll d;  
        cin >> d;  
        d--;  
  
        auto itr = st.find(PII(a[d], d));  
        PII pre, now, nxt;  
        if(itr != st.begin()) pre = *prev(itr);  
        now = *itr;  
        if(next(itr) != st.end()) nxt = *next(itr);  
  
        // treeからa[d]-pre(a[d]) と next(a[d])-a[d]を消す  
        if(itr != st.begin()) tree.erase_key(root, PII(now.first-pre.first, 0));  
        if(next(itr) != st.end()) tree.erase_key(root, PII(nxt.first-now.first, 0));  
        // treeにnext(a[d])-pre(a[d])を追加する  
        if(itr != st.begin() && next(itr) != st.end()) {  
            tree.insert_key(root, PII(nxt.first-pre.first, nxt.first-pre.first));  
        }  
        // stからPII(a[d], d)を消す  
        st.erase(PII(a[d], d));  
          
        cout << tree.query(root, 0, n-k-1-i).second << endl;  
    }  
  
    return 0;  
}