kmjp's blog

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

天下一プログラマーコンテスト2016 予選B : E - 天下一合体

このテクも勉強になった。
http://tenka1-2016-qualb.contest.atcoder.jp/tasks/tenka1_2016_qualB_e

問題

N要素の数列Aがある。
連続する2要素を選択し、結合して両者の和の値を取る1要素に置き換える、という処理を繰り返す。
Aの要素がM個になった段階で、各要素の絶対値の和を取る。
最適な順で要素を結合した場合、要素の絶対値の和を最小化せよ。

解法

この処理では結合した各要素のみが関係し、出来上がる要素は結合順には関係しない。
よって結局N要素の数列を以下に最適にM個に分割するかという問題になる。

先に累積和を取り、(A[i]を1-originとして) S[i]=A[i] + S[i-1]としておく。
計算量を考えなければ、以下のDPを行えば行えばよいことがわかる。
dp(i,j) := A[1...i]をj個の要素に分割したとき、各分割要素の和の絶対値の和の最小値
dp(i,j) = min_k(dp(k,j-1) + |S[i]-S[k]|)

ただ、上記DPはO(N^2*M)かかるので間にあわない。
上記minの内側は、S[i]<S[k]なら(dp(k,j-1) + S[k]-S[i])、そうでないなら(dp(k,j-1) + S[i]-S[k])である。
これらの最小値を高速に探す方法を考えよう。

それには、min(dp(k,j-1) + S[k])とmin(dp(k,j-1)-S[k])を求める2つのRMQがあればよい。
各kに対し、(dp(k,j-1)+S[k])と(dp(k,j-1)-S[k])のどちらを用いるかは、S[i]とS[k]の大小関係で決まる。
よって、RMQ中の添え字X[i]をS[i]の昇順で定めておこう。

すると以下の式になる。X[i]を境に2つのRMQをそれぞれ検索すればよくなり、計算量がO(NM logN)に収まる。
 \displaystyle dp(i,j) = min( min_{k=1 and X_k \lt X_i} (dp(k,j-1) - S_k) + S_i, min_{k=1 and X_k \gt X_i} (dp(k,j-1) + S_k) - S_i)

template<class V,int NV> class SegTree_1 {
public:
	vector<V> val;
	static V const def=1LL<<61;
	V comp(V l,V r){ return min(l,r);};
	
	SegTree_1(){val=vector<V>(NV*2,def);};
	V getval(int x,int y,int l=0,int r=NV,int k=1) { // x<=i<y
		if(r<=x || y<=l) return def;
		if(x<=l && r<=y) return val[k];
		return comp(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));
	}
	void update(int entry, V v) {
		entry += NV;
		val[entry]=v;
		while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]);
	}
};


SegTree_1<ll,1<<15> st[2];
int N,M;
ll S[202020];
pair<ll,int> P[202020];
int R[202020];
ll dp[20202][101];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	FOR(i,N) {
		cin>>x;
		S[i+1]=S[i]+x;
		P[i]={S[i+1],i};
	}
	sort(P,P+N);
	FOR(i,N) R[P[i].second+1]=i+1;
	
	memset(dp,0x3f,sizeof(dp));
	FOR(i,N) dp[i+1][1]=abs(S[i+1]);
	for(i=2;i<=M;i++) {
		FOR(x,st[0].val.size()) st[0].val[x]=st[1].val[x]=1LL<<61;
		
		for(x=1;x<=N;x++) {
			dp[x][i] = min(st[0].getval(0,R[x])+S[x],st[1].getval(R[x],N+1)-S[x]);
			st[0].update(R[x],dp[x][i-1]-S[x]);
			st[1].update(R[x],dp[x][i-1]+S[x]);
		}
	}
	
	cout<<dp[N][M]<<endl;
	
}

まとめ

DもEもRMQだな…。