全国統一プログラミング王決定戦予選 : F - Jewels

シンプルな問題設定だがわりとめんどい。
https://atcoder.jp/contests/nikkei2019-qual/tasks/nikkei2019_qual_f

問題

K通りの色の宝石が計N個ある。
各石の色C[i]と価値V[i]が与えられる。

N個中X個選ぶとき、「ある色の石が１個だけ選ばれる」という状態を認めないという条件のもとで、選んだ石の価値の総和の最大値を、X=1～Nのそれぞれについて求めよ。

解法

各色当然価値の高い順に選ぶのがいいので、事前に色毎に石を価値降順にソートしておこう。
なお、先頭２つは必ず同時に選ぶしかないので、事前に平均値に置き換えておく。

さて、全石を価値降順にソートし、先頭X個を選ぶとする。
もしX個目と(X+1)個目が同時に選ぶべきある石の最大価値の２個、という状況でないなら、単に先頭Xを選べばよい。
問題はそうでない場合である。ただ、この場合(X-1)個は必ず最大価値の(X-1)個を選んだ状態になっている。
そこで、(X-1)個からX個への遷移の仕方を考え、その最大値を選べばよい。

実際、Editorialによるの以下のいずれかが最大である。
3つめを思いつくのはなかなか大変そうだ。

• X個目をあきらめ、(X+2)個目以降で、1個だけ追加できる最大値の石を選ぶ。
• (X-1)個目以前で、1個だけ取り除ける最小価値の石を取り除き、X個目と(X+1)個目をセットで追加する。
• (X-1)個目以前で、2個セットで追加した石を取り除き、残された石のうち3個セットで追加できる石を追加する。

あとは、各状態の最小・最大の石ないし石のセットの集合を適宜更新していけばよい。

int N,K;
int C[202020];
ll V[202020];
vector<ll> cand[202020];
int id[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>C[i]>>V[i];
		C[i]--;
		cand[C[i]].push_back(V[i]*2);
	}
	
	multiset<pair<ll,int>> P,Q,R,S,T;
	vector<vector<ll>> ev;
	FOR(i,K) {;
		sort(ALL(cand[i]));
		cand[i][cand[i].size()-2]=cand[i][cand[i].size()-1]=(cand[i][cand[i].size()-1]+cand[i][cand[i].size()-2])/2;
		reverse(ALL(cand[i]));
		S.insert({cand[i][0]+cand[i][1],i});
		if(cand[i].size()>=3) T.insert({cand[i][0]+cand[i][1]+cand[i][2],i});
		FOR(j,cand[i].size()) ev.push_back({-cand[i][j],i,j});
	}
	
	sort(ALL(ev));
	ll cur=0;
	FOR(i,N) {
		auto e=ev[i];
		x=e[1];
		
		if(e[2]==0) {
			ll ret=-2;
			if(Q.size()) ret=max(ret,cur+Q.rbegin()->first);
			if(P.size()&&S.size()) ret=max(ret,cur-P.begin()->first+S.rbegin()->first);
			if(R.size()&&T.size()) ret=max(ret,cur-R.begin()->first+T.rbegin()->first);
			cout<<ret/2<<endl;
		}
		else {
			
			if(e[2]==1) {
				cur-=2*e[0];
				S.erase({cand[x][0]+cand[x][1],x});
				if(cand[x].size()>=3) T.erase({cand[x][0]+cand[x][1]+cand[x][2],x});
				R.insert({cand[x][0]+cand[x][1],x});
				id[x]=2;
				if(id[x]<cand[x].size()) Q.insert({cand[x][id[x]],x});
			}
			else {
				cur-=e[0];
				if(id[x]==2) R.erase({cand[x][0]+cand[x][1],x});
				else P.erase({cand[x][id[x]-1],x});
				Q.erase({cand[x][id[x]],x});
				P.insert({cand[x][id[x]],x});
				id[x]++;
				if(id[x]<cand[x].size()) Q.insert({cand[x][id[x]],x});
			}
			cout<<cur/2<<endl;
		}
		
	}
	
}

まとめ

「セットの集合」って変な書き方になってしまった。