Codeforces #413 E: Aquarium decoration

解説付きコード

/*
解説
二人とも好むアイテムの集合をX, MashaだけのはY, ArkadyだけのはZとする
全体でちょうどm個でなければならないが、とりあえずそのことは考えない。
Xからx個選ぶとする。xは|X|から0まで総当りする。
k-x個だけYとZから取る必要があるので取る。
ただし、YまたはZからk-x個取れない場合があるが、このとき条件を満たさない。
YとZから取った個数をy, zとする。ただし、y=z
明らかに
x+y+z<=m
である必要がある。
逆にこの時、残ったアイテムからm-(x+y+z)個を安い方から貪欲に取ればいい。
m-(x+y+z)個のアイテムの値段の和を得るために
Binary Indexed Treeを使う。
BIT を2個もちそれぞれcnt, sumとする。
アイテムが価格順であるとき、
cnt: クエリ[l, r)に対して[l, r)のうち残っているのはいくつあるか
sum: クエリ[l, r)に対して[l, r)のうち残っているものの値段の和
とする。
xは降順なので1個ずつcntとsumを調整できる。
また、xを降順に見ていったとき、y,zは昇順であることに注意すると
cnt,sumに対する更新はO(N)回で済む。
二分探索で
cntに対するクエリがm-(x+y+z)となるようなクエリ[0, p)を満たすpを求めれば、
m-(x+y+z)個の価格の和をsumに対してクエリ[0, p)を投げて求められる。
実行時間
O(n(log(n))^2)
*/
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int N, M, K;
    cin >> N >> M >> K;
    vector<pair<ll,int> > C(N);
    vi A(N), id(N);
    rep(i, N) {
        cin >> C[i].first;
        C[i].second = i;
    }
    // 値段は安い方を優先して使いたいので昇順ソート
    sort(all(C));
    vll cf(N);
    // ソート後のコスト、インデックス
    rep(i, N)cf[i] = C[i].first;
    rep(i, N)id[C[i].second] = i;
    for (int a : {1, 2}) {
        int t, x;
        cin >> t;
        while (t--) {
            cin >> x;
            --x;
            A[id[x]] |= a;
        }
    }
    vi X, Y, Z, DUMMY;
    rep(i, N) (A[i] == 3 ? X : A[i] == 2 ? Y : A[i] == 1 ? Z : DUMMY).push_back(i);
    // 集合X, Y, Zからいくつ選んだもののコストの和
    ll cur = 0;
    each(x, X)cur += cf[x];
    ll ans = LLONG_MAX;
    BIT cnt(N), sum(N);
    // Xに含まれないものをBITに加える
    rep(i, N)if (A[i] != 3) {
        cnt.add(i, 1);
        sum.add(i, cf[i]);
    }
    // 二分探索判定用
    int need = -1;
    auto f = [&](int m) {
        return cnt.sum(m) >= need;
    };
    auto g = [&](vi &W, int w, int a) {
        cur += a*cf[W[w]];
        cnt.add(W[w], -a);
        sum.add(W[w], -a*cf[W[w]]);
    };
    int y = 0, z = 0;
    for (int x = sz(X); x >= 0; --x) { // Xからx個選ぶ
        // YからK-x個以上選ばなければならない。
        while (x + y < K && y<sz(Y)) g(Y, y++, +1);
        // ZからK-x個以上選ばなければならない。
        while (x + z < K && z < sz(Z)) g(Z, z++, +1);
        // MashaかArkadyは不満足
        if (x + y < K || x + z < K)break;
        // x+y+z>Mとすると多すぎて駄目
        if (x + y + z <= M) {
            // あとM-x-y-z個必要
            need = M - x - y - z;
            if (need == 0)smin(ans, cur);
            else {
                int w = binarySearchR(0, N, f);
                smin(ans, cur + sum.sum(w));
            }
        }
        if (x) g(X, x - 1, -1);
    }
    cout << (ans != LLONG_MAX ? ans : -1) << endl;
}