kmjp's blog

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

Codeforces #162 Div2. E. Choosing Balls

これは少し考えて案が思いつかなかったので、解説を見て解いた問題。
http://codeforces.com/contest/265/problem/E

問題

最大100000個のボール列が与えられる。各ボールは、色と価値(整数)を持つ。
ここに、最大500個のクエリが来る。クエリにはAとBのパラメータが含まれている。
ボール列の部分列を作り、下記手順で得られるボール群の総スコアの最大値を返す問題。

  • 直前のボールと色が同じなら、このクエリにおけるこのボールのスコアはA*(ボールの価値)
  • 部分列の先頭または直前のボールと色が異なるなら、このクエリにおけるこのボールのスコアはB*(ボールの価値)

解法

A,Bやボールの価値が負の場合もあるので、余計なボールは部分列に含めないようにしなければならない。
ここはDPを用いて、直前のボールの色毎に最高スコアを取っておけばよい。
ただ問題は、Bが大きい場合など直前のボールが別の色の方が高いスコアになる場合、他の色のスコアの最大値を探さないといけない。
全色を合わせた最高スコアを覚えておく手もあるけど、たまたまその最高スコアを取ったときのボールの色が今のボールの色と一致してしまうと、結局その他の色の最高スコアを探さないといけない。

…結局全色の最高スコアを検索しなければならないか?と思ったら。
色別最高スコアは覚えておくんだけど、それとは別にスコア上位2色だけ覚えて置けばよかった。
最上位が今の色なら2位を見ればよい、というわけ。

そこで、ボールごとにスコア1位・2位の色を覚えながらDPしていけばO(N)で処理できる。
クエリ数と合わせてO(NQ)。2s近くかかったが元々上限5sの問題なので問題なし。

int N,Q;
int V[100001],C[100001];
ll MA;
ll CM[100001];
ll A,B;
int MC1,MC2;
void solve() {
	int f,r,i,j,k,l,cur,tar,ret,loop;
	int res,x,y;
	ll c;
	
	GET2(&N,&Q);
	FOR(i,N) V[i]=GETi();
	FOR(i,N) C[i]=GETi();
	
	FOR(i,Q) {
		A=GETi();
		B=GETi();
		fill(CM, CM+100001, -(1LL<<60));
		CM[0]=0;
		MA=MC1=MC2=0;
		FOR(x,N) {
			
			if(MC1==C[x]) c = max(CM[MC1] + V[x]*A, CM[MC2] + V[x]*B);
			else c = max(CM[MC1] + V[x]*B, CM[C[x]] + V[x]*A);
			
			if(c > CM[C[x]]) {
				CM[C[x]] = c;
				if(MC1!=C[x] && MC2!=C[x]) {
					if(c>=CM[MC1]) {
						MC2=MC1;
						MC1=C[x];
					}
					else if(c>=CM[MC2]) {
						MC2=C[x];
					}
				}
				else if(MC2==C[x]) {
					if(CM[MC2]>CM[MC1]) swap(MC1,MC2);
				}
			}
			MA=max(MA, c);
		}
		_P("%lld\n",MA);
	}
	
	return;
}

まとめ

2位まで覚えておけばよいということにぱっと考えが至らなかった…。