Codeforces #718 Div1+2 : E. Group Photo

なんか変な問題。
https://codeforces.com/contest/1517/problem/E

問題

N要素の整数列Aが与えられる。
これを2つの数列に分解しよう。
Aの添え字1～Nを、2つの数列C,Pに分けたとする。
その時、以下を満たす分け方は何通りか。

• C[i]-C[i-1]≦C[i+1]-C[i]
• P[i]-P[i-1]≧P[i+1]-P[i]
• Cに対応するAの要素の総和より、Pに対応するAの要素の総和の方が大きい。

解法

まず総和の条件は置いておく。
C[i]-C[i-1]＞2となるC[i]があったとすると、C[i-1]とC[i]の間に、P[j]-P[j-1]=1となる区間ができる。
これを踏まえると、あり得るパターンは正規表現で書くと以下のいずれかである。

• (P*)(C*)
• (P?)C*(CP)*P*(C?)

前者はPの数を総当たりすればよい。
後者は、最初のPと最後のCは2*2通り総当たりし、あとはC*の部分を総当たりしながら、尺取り法の要領で(CP*)を何通り取れるか計算すればよい。

int T,N;
int A[202020];
ll S[202002];
const ll mo=998244353;

ll slow() {
	int ret=0;
	int mask;
	int i;
	FOR(mask,1<<N) {
		ll sum=0;
		vector<int> B[2];
		FOR(i,N) {
			if(mask&(1<<i)) {
				B[1].push_back(i);
				sum+=A[i];
			}
			else {
				B[0].push_back(i);
				sum-=A[i];
			}
		}
		if(sum<=0) continue;
		if(B[0].size()==3&&B[0][1]-B[0][0]>B[0][2]-B[0][1]) continue;
		if(B[1].size()==3&&B[1][1]-B[1][0]<B[1][2]-B[1][1]) continue;
		ret++;
	}
	
	return ret;
}


ll hoge(vector<int> B, ll dif) {
	ll ret=0;
	int N=B.size();
	int i;
	FOR(i,N) S[i+1]=S[i]+B[i];
	
	ll T=S[N];
	for(int s=1;s<=2;s++) {
		ll CT=T-B[0];
		if(s==2) CT-=B[1];
	int L,R;
		for(L=R=s;L<N;L+=2) {
			while(R+2<=N-1&&(CT+dif-B[R+1]>T-(CT-B[R+1]))) CT-=B[R+1], R+=2;
			while(L<R&&(CT+dif<=T-(CT))) CT+=B[R-1], R-=2;
			if(CT+dif<=T-CT) break;
			ret+=((R-L)/2+1);
			CT-=B[L];
			if(L==R) R+=2, CT-=B[L+1];
		}
	}
	
	return ret;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		ll S=0;
		FOR(i,N) {
			cin>>A[i];
			S+=A[i];
		}
		
		if(N<=4) {
			cout<<slow()<<endl;
			continue;
		}
		ll ret=0;
		ll T=A[0];
		for(i=1;i<=N;i++) {
			if(T>S-T) ret++;
			T+=A[i];
		}
		FOR(x,2) FOR(y,2) {
			ll sum=0;
			vector<int> B;
			FOR(i,N) {
				if(i==0&&x==1) sum+=A[i];
				else if(i==N-1&&y==1) sum-=A[i];
				else B.push_back(A[i]);
			}
			ret+=hoge(B,sum);
		}
		cout<<ret%mo<<endl;
		
	}
}

まとめ

これ系苦手なんだけど本番通せてよかったな。