ABC396 - C - Buy Balls
2025/03/08に行われたAtCoder Beginner Contest 396のC問題の解法についての記事です。
考えたこと
価値が高い方から黒・白の玉双方を選ぶため、まずは黒の玉のvector、白の玉のvectorをそれぞれ降順にsortしてやることから始めます。
さて、白の玉を\( n \)個選んだとき、黒の玉の選び方は、
1. 大きい方から順に\( n \)個選ぶ
2. \( n + ( 1 , 2 , ... , k ) \)個目が負の数でないなら追加で選ぶ
3. \( n + k + 1 \)個目が初めて負の数になるなら、\( n + k \)個目までで選ぶのを止める
4. 黒の玉に負の数が含まれないなら最大の\( N \)個まで選ぶ
という操作をすることで、選んだ白と黒の玉の数の和を最大化できます。
あとは白の玉を\( n = ( 0 , 1 , 2 , ... , M ) \) 個選んだときの和を求め、最大値を導出すればよいです。
ただし、愚直に実装すると\( O ( N M ) \)時間かかってしまうため、今回の制限では制限時間に間に合いません。
あらかじめ黒の玉の和が最大になる選び方を計算しておくことで、\( O ( N + M ) \)時間で回答することができます。
提出したコード
#include <bits/stdc++.h>using namespace std;int main() {int B, W;cin >> B >> W;vector<long long> vecB(B);vector<long long> vecW(W);for (int i = 0; i < B; i++) {cin >> vecB.at(i);}for (int i = 0; i < W; i++) {cin >> vecW.at(i);}sort(vecB.rbegin(), vecB.rend());vector<long long> sumB(B);int Bminus = B;for (int i = 0; i < B; i++) {if (Bminus == B && vecB.at(i) < 0) {Bminus = i;}if (i == 0) {sumB.at(i) = vecB.at(i);}else {sumB.at(i) = sumB.at(i - 1) + vecB.at(i);}}sort(vecW.rbegin(), vecW.rend());vector<long long> sumW(W);for (int i = 0; i < W; i++) {if (i == 0) {sumW.at(i) = vecW.at(i);}else {sumW.at(i) = sumW.at(i - 1) + vecW.at(i);}}long long ans = -10000000000000000;for (int i = -1; i < W && i < B; i++) {int j = i;if (j < Bminus - 1) {j = Bminus - 1;}if (j >= B) {j = B - 1;}if (j < 0) {j = 0;}if (j < i) {continue;}long long num;if (i == -1) {num = sumB.at(j);} else {num = sumW.at(i) + sumB.at(j);}ans = max(num, ans);}if (ans < 0) {ans = 0;}cout << ans << endl;return 0;}
解法
黒い玉の数字の列vecBと、白い玉の数字の列vecWを用意して、それぞれに数値を代入したあと、和を求める処理を行っていきます。
ここでは先に黒い玉の処理を行ったあとで、白い玉の方を処理しています。黒い玉では和の数列が広義単調増加から初めて減少に転じるときの玉の個数Bminusも求めておきます。
あとは白い玉の個数を増やしていきながら、黒い玉との総和が最大となるとき、つまりそのときの総和＝答えを求めていきます。ここで、白い玉が全て負の数、かつ黒い玉に正整数が含まれる場合に、白い玉を0個選ぶ必要があるため、白い玉の数を表す\( i \)を\( -1 \)から始めることで対処しました。うまくプログラムすれば\( i == 0 \) && \( j == 0 \)で白黒ともに0個のときを表すことができて、直感的に分かりやすく組むことができると思います。
もう少し綺麗な書き方があると思いますが、これでも回答となります。