kmjp's blog

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

Codeforces #344 Div2. E. Product Sum

うーん、あと一歩。
http://codeforces.com/contest/631/problem/E

問題

N要素の数列A[1...N]が与えられる。
数列のスコアCを \displaystyle C=\sum_{i=1}^N A_i*iで定義する。

数列中、1つの要素を移動し、任意の場所に差し替えることができるとき、スコアの最大値を求めよ。

解法

まず初期状態のスコアを求めよう。

要素を1つ移動するということは、数列の連続部分列を左右どちらかにローテートすることに相当する。
ローテートしたとき、スコアの差分を考えよう。この差分の最大値が求められれば、この問題は解けたことになる。

S[i] := sum(A[1...i])とする。
まず、A[l...r]を左ローテートしたときの差分DL(l,r) (すなわちA[l]をr番目となるよう回転)を考える。
DL(l,r) = A[l]*(r-l)-S[r]+S[l] = A[l]*r - S[r] + S[l] - A[l]*l
なので、f(r,x) = r*x - S[r]と置くと、DL(l,r) = f(r,A[l]) + S[l] - Aと置ける。
lを固定すると、DL(l,*)の最大値は結局f(r,A[l])の最大ととなる。
これはConvex Hull Trickの要領でxをA[l]に固定したときのf(r,x)の最大値を二分探索で求められるようにしておくと、f(r,A[l])の最大値を求められる。

同様にA[l...r]を右ローテートした場合、rを固定してDR(l,r) = g(l,A[r]) + S[r-1] - A[r]*r, g(l,x) = l*x - S[l-1]とすると、同様にgに対する最大値を二分探索で求められる。
fとgはl,rを動かす方向が逆なので、Convex Hull Trickの過程で直線を残す・削除する大小基準が反転することに注意しよう。
そうしないとテストケース21,48あたりで引っかかる。

なお、この問題はテストケースが弱くもっとお手軽な二分探索でも通ってしまうようだ。

template<typename V> struct ConvexHull {
	deque<pair<V,V>> Q;
	int cmptype=1; // 0-min 1-max
	V calc(pair<V,V> p, V x) {
		return p.first*x+p.second;
	}
	int dodo(pair<V,V> A,pair<V,V> B, pair<V,V> C) { // max or min
		return cmptype^((B.second-C.second)*(B.first-A.first)<=(A.second-B.second)*(C.first-B.first));
	}
	void add(V a, V b) { // add ax+b
		Q.push_back({a,b});
		int v;
		while((v=Q.size())>=3 && dodo(Q[v-3],Q[v-2],Q[v-1]))
			Q[v-2]=Q[v-1], Q.pop_back();
	}
	
	V query(V x) {
		int L=-1,R=Q.size()-1;
		while(R-L>1) {
			int M=(L+R)/2;
			((calc(Q[M],x)<=calc(Q[M+1],x))?L:R)=M;
		}
		return calc(Q[R],x);
	}
};

int N;
ll A[202020];
ll S[202020];
ll ma;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i+1];
		ma+=(i+1LL)*A[i+1];
		S[i+1]=S[i]+A[i+1];
	}
	
	ConvexHull<ll> ch;
	ll dif=0;
	ll L,R;
	// left shift
	for(L=N;L>=1;L--) {
		ch.add(L,-S[L]);
		dif=max(dif,ch.query(A[L])+S[L]-A[L]*L);
	}
	// right shift
	ch.Q.clear();
	ch.cmptype=0;
	for(R=1;R<=N;R++) {
		ch.add(R,-S[R-1]);
		dif=max(dif,ch.query(A[R])+S[R-1]-A[R]*R);
	}
	
	cout<<ma+dif<<endl;
}

まとめ

本番、Convex Hull Trick使いそうだなーとは思ったものの、うまく式をテクが使える形に変形できなかった。
…うまく行ったとしても、上に書いた大小判断のせいで落としてそう。