[AtCoder] ABC 160 E – Red and Green Apples
問題
https://atcoder.jp/contests/abc160/tasks/abc160_e方針
全探索
赤色、緑色、無色のリンゴの美味しさを降順に並び替え、その累積和をそれぞれ、\( P, Q, R \) とします。また、\( i \) 個のリンゴの累積和を \( P(i) \) のように表します。次に、無色のリンゴを \( 0 \) 個から \( \min(C, X + Y)\) 個まで食べる時の美味しさの総和を求めます。無色のリンゴを食べた個数を \( z \) とします。無色のリンゴを\( z \) 個食べたとき、総和が最大となるような赤色と緑色のリンゴの個数をそれぞれ \( X_z, Y_z \) とします。
\( z = 0 \) のとき
このとき、\( X_0 = X , Y_0 = Y\) となるので、\( P(X) + Q(Y) \) が総和の最大です。
\( z = 1 \) のとき
このとき、
\[ R(1) + \max(P(X_0 – 1) + Q(Y), P(X) + Q(Y_0 – 1) )\]
が総和の最大です。
\( z = 2 \) のとき
このとき、\( z = 1 \) の状態から遷移するので、
\[ R(2) + \max(P(X_1 – 1) + Q(Y), P(X) + Q(Y_1 – 1) )\]
が総和の最大です。
\( z = i \) のとき
このとき、
\[ R(i) + \max(P(X_{i-1} – 1) + Q(Y), P(X) + Q(Y_{i-1} – 1) )\]
が総和の最大です。よって、全探索によって求めることができます。
整列
赤色のリンゴは最大でも \( X \) 個食べれば十分であり、緑色のリンゴは最大でも \( Y \) 個食べれば十分です。したがって、美味しさの高い順にそれぞれ \( X , Y\) 個ずつ選び、無色のリンゴと合わせて、 \( X + Y + C \) 個の中から美味しさの高い順に \( X + Y \) 個食べればいいです。
コード
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int main() {
int X, Y, A, B, C;
cin >> X >> Y >> A >> B >> C;
vector<ll> p(A), q(B), r(C);
for (int i = 0; i < A; i++) {
cin >> p[i];
}
for (int i = 0; i < B; i++) {
cin >> q[i];
}
for (int i = 0; i < C; i++) {
cin >> r[i];
}
sort(p.begin(), p.end(), greater<int>());
sort(q.begin(), q.end(), greater<int>());
sort(r.begin(), r.end(), greater<int>());
ll P[A + 1]{}, Q[B + 1]{}, R[C + 1]{};
for (int i = 0; i < A; i++) {
P[i + 1] = P[i] + p[i];
}
for (int i = 0; i < B; i++) {
Q[i + 1] = Q[i] + q[i];
}
for (int i = 0; i < C; i++) {
R[i + 1] = R[i] + r[i];
}
ll ans = P[X] + Q[Y];
int id_p = X;
int id_q = Y;
for (int i = 1; i <= min(C, X + Y); i++) {
if (id_p == 0 && id_q == 0) {
ans = max(ans, R[min(X + Y, C)]);
break;
}
if (id_p == 0) {
id_q--;
} else if (id_q == 0) {
id_p--;
} else if (P[id_p - 1] + Q[id_q] < P[id_p] + Q[id_q - 1]) {
id_q--;
} else {
id_p--;
}
ans = max(ans, R[i] + P[id_p] + Q[id_q]);
}
cout << ans << "\n";
return 0;
}