kmjp's blog

競技プログラミング参加記です

AtCoder ABC #218 : H - Red and Blue Lamps

これは解いておきたかったな。
https://atcoder.jp/contests/abc218/tasks/abc218_h

問題

N個のランプがあるので、R個を選んで赤、残り(N-R)個を青くさせることを考える。
i番とi+1番の色が異なる場合、スコアA[i]を得られるとする。

任意の点灯のさせ方における、総スコアの最大値はいくつか。

解法

以下赤の方が半分以下と仮定する(半分を超えるなら、赤と青を逆にする)。
この時、赤色が2個以上連続する意味はない。

さて、N個中赤くする箇所をR個選ぶことを考える。
赤が連続しないことを考えると、i番を赤くするとA[i]+A[i-1]のスコアが増えると考えることができる。(A[N]=0とすれば、N番目を選択した際も不具合が生じない)。

実は、上記条件は以下のように言い換えることができる。

  • いくつかスコアに計上されたランプがわかっているとき、新たにi番を赤くするとA[i]+A[j] (jはi未満のうち、スコア未計上のタンプ番号の最大値)だけスコアが増える。

なぜなら、(j+1)~(i-1)番のランプがスコア計上済みということは、(i-1)、(i-3)、…、(j+2)番のランプが赤くなっていることを意味する。
ここを、i、(i-2)、…、(j+1)番のランプを赤くすることにすれば、赤くするランプは1つ増えるだけで、i番とj番をスコアに計上できる形にすることができる。

あとは、上記手順をシミュレートしながら、A[i]+A[j]の最大値をどん欲にR個選ぼう。
(i番とj番がスコアに計上されることで、以後j番を赤色として選択できなくなるし、i番の次の未計上のランプのスコアも変化することに留意。)

int N,R;
ll A[202020];
ll B[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>R;
	FOR(i,N-1) cin>>A[i+1];
	set<int> alive;
	set<pair<int,int>> cand;
	
	for(i=1;i<=N;i++) {
		B[i]=A[i]+A[i-1];
		alive.insert(i);
		cand.insert({-B[i],i});
	}
	
	if(R>N-R) R=N-R;
	ll ret=0;
	while(R--) {
		auto p=*cand.begin();
		cand.erase(cand.begin());
		ret+=-p.first;
		x=p.second;
		alive.erase(x);
		auto it=alive.lower_bound(x);
		if(it!=alive.begin()) {
			y=*prev(it);
			if(y>0) {
				cand.erase({-B[y],y});
				alive.erase(y);
			}
		}
		if(it!=alive.end()) {
			x=*it;
			cand.erase({-B[x],x});
			B[x]=A[x];
			auto it=alive.lower_bound(x);
			if(it!=alive.begin()) {
				B[x]+=A[*prev(it)];
			}
			cand.insert({-B[x],x});
		}
	}
	cout<<ret<<endl;
	
}

まとめ

Candiesのこと頭をよぎったけど、問題名を覚えておらず探せなかった…。