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; }