Codeforces 162C (264C) Choosing Balls

問題

n個のボールがあり、それぞれには値v[i]と色c[i]がつけられている。
今、n個のボールから順番を保存したまま何個かのボールを抜き出して列を作る。


この列に対して、

  • 列の最初以外かつ、直前のボールと同じ色のボールが来るたびにa * v[i]点
  • それ以外の場合のボールが来るたびにb * v[i]点

が得られるとする。

a, bがq通り与えられるとき、それぞれのa, bについて、
得られる最大の得点を求めよ。

制約条件

n≦10^5

v ≦10^5

1≦c≦n
q≦500

方針

dp[i] := 最後に色iのボールを使ったときの得点の最大値
とする。
これを更新するようなDPをしたいが、新しいボール以外の色の最大値を短い時間で調べたい。


これには、最大値を大きい順に2つ持っておけば、
どちらかが今のボールとは違う色の最大値になっている。
したがって、全体でO(n)のDPで得られる最大の得点が計算できる。

ソースコード

int n, q;
int v[MX], c[MX];
ll dp[MX];

int main(){
	cin >> n >> q;
	rep(i, n) cin >> v[i];
	rep(i, n) cin >> c[i];
	rep(it, q){
		rep(i, n + 1) dp[i] = -inf;
		
		int col[2] = {0, 0};
		ll mx[2] = {-inf, -inf};
		ll ans = 0;
		ll a, b;
		
		cin >> a >> b;
		
		rep(i, n){
			ll nxt = max(dp[c[i]] + a * v[i], b * v[i]);
			rep(j, 2) nxt = max(nxt, mx[j] + (col[j] == c[i] ? a : b) * v[i]);
			dp[c[i]] = max(dp[c[i]], nxt);
			ans = max(ans, nxt);
			
			int idx = -1;
			rep(j, 2) if(col[j] == c[i]) idx = j;
			if(idx >= 0){
				if(mx[idx] < nxt) mx[idx] = nxt;
			}
			else{
				idx = mx[0] < mx[1] ? 0 : 1;
				if(mx[idx] < nxt) mx[idx] = nxt, col[idx] = c[i];
			}
		}
		cout << ans << endl;
	}
	return 0;
}