Playrix Codescapes Cup (Codeforces Round #413, rated, Div. 1 + Div. 2) E - Aquarium decoration 贪心 + 平衡树

E - Aquarium decoration

枚举两个人都喜欢的个数，就能得到单个喜欢的个数，然后用平衡树维护前k大的和。

#include<bits/stdc++.h>
#define LL long long
#define fi first
#define se second
#define mk make_pair
#define PII pair<int, int>
#define PLI pair<LL, int>
#define ull unsigned long long
using namespace std;

const int N = 2e5 + ;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = ;
const int Mod = 1e9 + ;

int n, m, k, num, cnt, state[N], d[N], a[N], b[N], c[N];
LL suma[N], sumb[N], sumc[N];

struct node {
    node* ch[];
    int key, fix, sz, cnt;
    LL sum;
    void update() {
        sz = ch[]->sz + ch[]->sz + cnt;
        sum = ch[]->sum + ch[]->sum + 1ll*cnt*key;
    }
};

typedef node* P_node;

struct Treap {
    node base[N], nil;
    P_node root, null, len;
    Treap() {
        root = null = &nil;
        null->key = null->fix = 1e9;
        null->sz = null->cnt = ;
        null->ch[] = null->ch[] = null;
        len = base;
    }
    P_node newnode(int tkey) {
        len->key = tkey;
        len->fix = rand();
        len->ch[] = len->ch[] = null;
        len->sz = len->cnt = ;
        len->sum = tkey;
        return len++;
    }
    void rot(P_node &p, int d) {
        P_node k = p->ch[d ^ ];
        p->ch[d ^ ] = k->ch[d];
        k->ch[d] = p;
        p->update();
        k->update();
        p = k;
    }
    void _Insert(P_node &p, int tkey) {
        if(p == null) {
            p = newnode(tkey);
        } else if(p->key == tkey) {
            p->cnt++;
        } else {
            int d = tkey > p->key;
            _Insert(p->ch[d], tkey);
            if(p->ch[d]->fix > p->fix) {
                rot(p, d ^ );
            }
        }
        p->update();
    }

    void _Delete(P_node &p, int tkey) {
        if(p == null) return;
        if(p->key == tkey) {
            if(p->cnt > ) p->cnt--;
            else if(p->ch[] == null) p = p->ch[];
            else if(p->ch[] == null) p = p->ch[];
            else {
                int d = p->ch[]->fix > p->ch[]->fix;
                rot(p, d);
                _Delete(p->ch[d], tkey);
            }
        } else {
            _Delete(p->ch[tkey > p->key], tkey);
        }
        p->update();
    }
    int _Kth(P_node p, int k) {
        if(p == null || k <  || k > p->sz) return ;
        if(k < p->ch[]->sz + ) return _Kth(p->ch[], k);
        if(k > p->ch[]->sz + p->cnt) return _Kth(p->ch[], k - p->ch[]->sz - p->cnt);
        return p->key;
    }
    int _Rank(P_node p, int tkey, int res) {
        if(p == null) return -;
        if(p->key == tkey) return p->ch[]->sz + res + ;
        if(tkey < p->key) return _Rank(p->ch[], tkey, res);
        return _Rank(p->ch[], tkey, res + p->ch[]->sz + p->cnt);
    }
    int _Pred(P_node p, int tkey){
        if(p == null) return -1e9;
        if(tkey <= p->key) return _Pred(p->ch[], tkey);
        return max(p->key, _Pred(p->ch[], tkey));
    }
    int _Succ(P_node p, int tkey){
        if(p == null) return 1e9;
        if(tkey >= p->key) return _Succ(p->ch[], tkey);
        return min(p->key, _Succ(p->ch[], tkey));
    }
    LL _Query(P_node p, int res) {
        if(!res) return ;
        if(p->ch[]->sz >= res) return _Query(p->ch[], res);
        else if(p->ch[]->sz + p->cnt < res) {
            return p->ch[]->sum + 1ll*p->key*p->cnt + _Query(p->ch[], res - p->ch[]->sz - p->cnt);
        } else {
            return p->ch[]->sum + 1ll*p->key*(res - p->ch[]->sz);
        }
    }
    void Insert(int tkey){ _Insert(root,tkey); }
    void Delete(int tkey){ _Delete(root,tkey); }
    int Kth(int k){ return _Kth(root,k); }
    int Rank(int tkey){ return _Rank(root,tkey,); }
    int Pred(int tkey){ return _Pred(root,tkey); }
    int Succ(int tkey){ return _Succ(root,tkey); }
    LL Query(int res){ return _Query(root,res); }
}tp;

int main() {
    scanf("%d%d%d", &n, &m, &k);
    for(int i = ; i <= n; i++) scanf("%d", &d[i]);
    scanf("%d", &num);
    for(int i = ; i <= num; i++) {
        int x; scanf("%d", &x);
        state[x] |= ;
    }
    scanf("%d", &num);
    for(int i = ; i <= num; i++) {
        int x; scanf("%d", &x);
        state[x] |= ;
    }
    for(int i = ; i <= n; i++) {
        if(state[i] == ) tp.Insert(d[i]), cnt++;
        else if(state[i] == ) a[++a[]] = d[i];
        else if(state[i] == ) b[++b[]] = d[i];
        else c[++c[]] = d[i];
    }
    sort(a + , a +  + a[]);
    sort(b + , b +  + b[]);
    sort(c + , c +  + c[]);
    for(int i = ; i <= a[]; i++)
        suma[i] = suma[i-] + a[i];
    for(int i = ; i <= b[]; i++)
        sumb[i] = sumb[i-] + b[i];
    for(int i = ; i <= c[]; i++)
        sumc[i] = sumc[i-] + c[i];

    LL ans = INF;
    for(int i = ; i <= c[]; i++) tp.Insert(c[i]), cnt++;
    for(int i = ; i <= c[]; i++) {
        if(i) tp.Delete(c[i]), cnt--;
        int res1 = max(, k - i);
        if(a[] < res1 || b[] < res1) continue;
        while(a[] > res1) {
            tp.Insert(a[a[]]);
            a[]--; cnt++;
        }
        while(b[] > res1) {
            tp.Insert(b[b[]]);
            b[]--; cnt++;
        }
        if(i +  * res1 > m) continue;
        int res2 = m - i -  * res1;
        ans = min(ans, sumc[i] + suma[a[]] + sumb[b[]] + tp.Query(res2));
    }
    printf("%lld\n", ans == INF ? - : ans);
    return ;
}

/*
*/