CF799E Aquarium decoration

我们将物品分成 \(4\) 个集合：

• 集合 \(A\) 包含只有 Arkady 喜欢的装饰。
• 集合 \(B\) 包含只有 Masha 喜欢的装饰。
• 集合 \(C\) 包含二人都喜欢的装饰。
• 集合 \(D\) 包含没人喜欢的装饰。

分别记这 \(4\) 个集合的大小为 \(a,b,c,d\)。

我们可以先将这 \(4\) 个集合分别按照花费从小到大来排序，并在后面按照这样的顺序来选择装饰。

不妨让我们来枚举我们从集合 \(C\) 中选取了多少个装饰，假设选择了 \(i\) 个。

那么我们至少要从集合 \(A\) 和集合 \(B\) 中分别选取 \(j=k-i\) 个装饰。

那么对于确定的 \(i\)，这 \(i+2j\) 个装饰是必选的，我们先将它们选上。此时还有一个限制没有满足，就是 \(i+2j\) 可能不足 \(m\)。

此时，我们就需要从集合 \(A\) 的剩下 \(a-i\)，集合 \(B\) 的剩下 \(b-j\)，集合 \(C\) 的剩下 \(c-j\) 和集合 \(D\) 的 \(d\) 个装饰中选花费最小的 \(m-i-2j\) 个装饰。

考虑用一棵权值线段树来维护。我们将剩下的装饰以花费作为下标添加的线段树上，并维护区间的花费和以及拥有装饰的数量。

每次相当于在线段树上二分找到第一个满足装饰数量前缀和大于等于 \(m-i-2j\) 的节点。

时间复杂度 \(O(n\log V)\)，可以通过离散化做到 \(O(n\log n)\)。

#include<bits/stdc++.h>
typedef long long ll;
using namespace std;
const int N = 2e5 + 5, M = 1e9;
int n, m, qwq, K, c[N], a[N], b[N];
struct node {
    int t, id;
    bool operator <(const node &p) const {return t < p.t;}
} arr[N], brr[N], crr[N];
int cnta, cntb, cntc;
ll suma[N], sumb[N], sumc[N], ans = 0x3f3f3f3f3f3f3f3f;
int rot, siz[N << 5], lc[N << 5], rc[N << 5], cnt;;
ll tr[N << 5];
void change(int &p, int l, int r, ll x, int opt) {
    if(!p) p = ++cnt;
    int mid = l + r >> 1;
    if(l == r) return tr[p] += x * opt, siz[p] += opt, void();
    if(x <= mid) change(lc[p], l, mid, x, opt);
    else if(mid < x) change(rc[p], mid + 1, r, x, opt);
    tr[p] = tr[lc[p]] + tr[rc[p]], siz[p] = siz[lc[p]] + siz[rc[p]];
}
ll query(int p, int l, int r, int k) {
    int mid = l + r >> 1;
    if(l == r) return 1ll * k * l;
    if(siz[lc[p]] >= k) return query(lc[p], l, mid, k);
    else return tr[lc[p]] + query(rc[p], mid + 1, r, k - siz[lc[p]]);
}
int main() {
#ifdef ddxrS
    freopen("sample.in", "r", stdin);
    freopen("sample.out", "w", stdout);
#endif
    cin>>n>>m>>K;
    for(int i = 1; i <= n; i++) cin>>c[i];
    cin>>qwq; for(int i = 1, p; i <= qwq; i++) cin>>p, a[p] = true;
    cin>>qwq; for(int i = 1, p; i <= qwq; i++) cin>>p, b[p] = true;
    for(int i = 1; i <= n; i++) {
        if(a[i] && b[i]) crr[++cntc] = {c[i], i};
        else if(a[i]) arr[++cnta] = {c[i], i};
        else if(b[i]) brr[++cntb] = {c[i], i};
        else change(rot, 1, M, c[i], 1);
    }
    sort(arr + 1, arr + cnta + 1), sort(brr + 1, brr + cntb + 1), sort(crr + 1, crr + cntc + 1);
    for(int i = 1; i <= cnta; i++) suma[i] = suma[i - 1] + arr[i].t;
    for(int i = 1; i <= cntb; i++) sumb[i] = sumb[i - 1] + brr[i].t;
    for(int i = 1; i <= cntc; i++) sumc[i] = sumc[i - 1] + crr[i].t;
    for(int i = K + 1; i <= cnta; i++) change(rot, 1, M, arr[i].t, 1);
    for(int i = K + 1; i <= cntb; i++) change(rot, 1, M, brr[i].t, 1);
    for(int i = 1; i <= cntc; i++) change(rot, 1, M, crr[i].t, 1);
    for(int i = 0, j = min(cnta, K), k = min(cntb, K); i <= cntc && i <= K; i++) {
        while(j >= 1 && j > K - i) change(rot, 1, M, arr[j].t, 1), j--;
        while(k >= 1 && k > K - i) change(rot, 1, M, brr[k].t, 1), k--;
        if(i >= 1) change(rot, 1, M, crr[i].t, -1);
        if(j >= K - i && k >= K - i && i + j + k <= m && siz[rot] >= m - i - j - k) ans = min(ans, suma[j] + sumb[k] + sumc[i] + query(rot, 1, M, m - i - j - k));
    }
    cout<<(ans < 0x3f3f3f3f3f3f3f3f ? ans : -1)<<'\n';
    cerr<<clock() * 1.0 / CLOCKS_PER_SEC<<'\n';
    return 0;
}