kmjp's blog

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

Codeforces #783 : Div1 B. Optimal Partition

出来が悪くはなかった回。
https://codeforces.com/contest/1667/problem/B

問題

整数列Aが与えられる。
これらをいくつかの連続部分列に分割したとする。
各連続部分列のスコアは、

  • 総和が正なら、要素数分
  • 総和が0なら、0
  • 総和が負なら、-(要素数分)

になるとする。最適な分割をした時、総スコアの最大値を答えよ。

解法

Aのprefix sumを取り、昇順となるようindexを振りなおそう。
SegTreeに、Aのprefixにおける(スコア-index)の最大値を載せておき、区間最大値を求められるようにすれば、Aのprefixをだんだん長くしながら、(スコア-index)の最大値を順次求めることができる。

int T,N;
ll A[505050],S[505050];

template<class V,int NV> class SegTree_1 {
public:
	vector<V> val;
	static V const def=-(1<<25);
	V comp(V l,V r){ return max(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]=comp(v,val[entry]); //上書きかチェック
		while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]);
	}
	void update2(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<int,1<<20> st;
int same[505050];
int dp[505050];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		vector<ll> cand;
		cand.push_back(0);
		S[0]=0;
		FOR(i,N) {
			cin>>A[i];
			S[i+1]=S[i]+A[i];
			cand.push_back(S[i+1]);
		}
		sort(ALL(cand));
		FOR(i,N+1) {
			st.update2(i,-1<<25);
			S[i]=lower_bound(ALL(cand),S[i])-cand.begin();
			same[i]=-1<<25;
		}
		st.update2(N+1,-1<<30);
		
		st.update(S[0],0);
		same[S[0]]=0;
		dp[0]=0;
		ll ret;
		for(i=1;i<=N;i++) {
			dp[i]=max(same[S[i]],st.getval(0,S[i])+i);
			if(A[i-1]==0) dp[i]=max(dp[i],dp[i-1]);
			if(A[i-1]<0) dp[i]=max(dp[i],dp[i-1]-1);
			if(A[i-1]>0) dp[i]=max(dp[i],dp[i-1]+1);
			if(i==N) {
				ret=dp[i];
			}
			else {
				same[S[i]]=max(same[S[i]],dp[i]);
				st.update(S[i],dp[i]-i);
			}
		}
		
		cout<<ret<<endl;
	}
}

まとめ

本番割と手間取ってるな。