Codeforces 799E Aquarium decoration

考虑去枚举能满足 \(1, 2\) 的个数 \(l\)，那自然只能满足 \(1\) 或 \(2\) 的个数为 \(\max\{k - l, 0\}\)。
对于还剩下的，可以从只能满足 \(1, 2\) 和不能满足任意一个的选出来。
显然如果要选就是选最小的。

考虑用个数据结构优化这个过程。
查询前 \(k\) 大的和，插入一个数，可以想到用值域线段树。
对于一开始就可以把都不能满足的插入线段树。
\(k\) 递增，\(\max\{k - l, 0\}\) 递减，于是可以依次取出只能满足 \(1\) 或 \(2\) 两个序列中分别的最大值，然后放入线段树中并从序列中删除，因为其肯定不可能再在只能满足 \(1\) 个的里面出现了。

注意一些边界情况。

时间复杂度 \(O(n\log V)\)。

#include<bits/stdc++.h>
using i64 = long long;
const i64 inf = LLONG_MAX;
const int maxn = 2e5 + 10, maxp = maxn * 31;
i64 sum[maxp]; int cnt[maxp], ls[maxp], rs[maxp], tot, rt;
void add(int &k, int x, int l = 1, int r = 1e9) {
    if (! k) k = ++tot;
    cnt[k]++, sum[k] += x;
    if (l == r) return ;
    int mid = (l + r) >> 1;
    x <= mid ? add(ls[k], x, l, mid) : add(rs[k], x, mid + 1, r);
}
i64 qry(int k, int x, int l = 1, int r = 1e9) {
    if (l == r) return 1ll * l * x;
    int mid = (l + r) >> 1;
    return cnt[ls[k]] < x ? (sum[ls[k]] + qry(rs[k], x - cnt[ls[k]], mid + 1, r)) : qry(ls[k], x, l, mid);
}
std::vector<int> a[4];
int A[maxn], b[maxn];
int main() {
    int n, m, k; scanf("%d%d%d", &n, &m, &k);
    for (int i = 1; i <= n; i++) scanf("%d", &A[i]);
    int p; scanf("%d", &p);
    while (p--) {int x; scanf("%d", &x); b[x] |= 1;}
    scanf("%d", &p);
    while (p--) {int x; scanf("%d", &x); b[x] |= 2;}
    for (int i = 1; i <= n; i++) a[b[i]].push_back(A[i]);
    for (int i = 0; i < 4; i++) std::sort(a[i].begin(), a[i].end());
    for (int x : a[0]) add(rt, x);
    int l = std::max({0, k - int(a[1].size()), k - int(a[2].size())});
    i64 sum1 = 0, sum2 = 0;
    for (i64 x : a[1]) sum1 += x;
    for (i64 x : a[2]) sum2 += x;
    i64 ans = inf, sum3 = 0;
    if (l > a[3].size()) return puts("-1"), 0;
    for (int i = 0; i < l; i++) sum3 += a[3][i];
    while (l <= m && l <= a[3].size()) {
        int p = std::max(k - l, 0);
        while (a[1].size() > p) add(rt, a[1].back()), sum1 -= a[1].back(), a[1].pop_back();
        while (a[2].size() > p) add(rt, a[2].back()), sum2 -= a[2].back(), a[2].pop_back();
        if (l + p + p <= m) ans = std::min(ans, sum1 + sum2 + sum3 + qry(rt, m - l - p - p));
        if (l == a[3].size()) break;
        sum3 += a[3][l++];
    }
    printf("%lld\n", ans == inf ? -1 : ans);
    return 0;
}