AtCoder ARC #153 : C - ± Increasing Sequence

ちょっと微妙な出来だった回。
https://atcoder.jp/contests/arc153/tasks/arc153_c

問題

1と-1で構成されるN要素の整数列Aが与えられる。
単調増加なN要素の整数列Xのうち、AとXの積和が0となるXが存在するならそれを答えよ。

解法

X={0,1,2,...}としたとき、AとXの積和Sが正だったとする。(負の場合も同じように考えられる。)

AのM要素のprefix sumが1だったとすると、Xの先頭M要素をデクリメントすればSを1減らせる。
または、AのM要素のsuffix sumが-1だったとすると、Xの末尾M要素をインクリメントすればSを1減らせる。

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

void out() {
	cout<<"Yes"<<endl;
	int i;
	ll sum=0;
	FOR(i,N) {
		cout<<X[i]<<" ";
		sum+=X[i]*A[i];
	}
	cout<<endl;
	assert(sum==0);
	exit(0);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	ll sum=0;
	FOR(i,N) {
		cin>>A[i];
		S[i+1]=S[i]+A[i];
		X[i]=i;
		sum+=A[i]*X[i];
	}
	
	if(sum==0) out();
	
	if(sum>0) {
		int cur=0;
		FOR(i,N) {
			cur+=A[i];
			if(cur==1) {
				ll s=sum;
				FOR(j,i+1) {
					X[j]-=s;
					sum-=s*A[j];
				}
				out();
				break;
			}
		}
		cur=0;
		for(i=N-1;i>=0;i--) {
			cur+=A[i];
			if(cur==-1) {
				ll s=sum;
				for(j=i;j<N;j++) {
					X[j]+=s;
					sum+=s*A[j];
				}
				out();
				break;
			}
		}
	}
	if(sum<0) {
		int cur=0;
		FOR(i,N) {
			cur+=A[i];
			if(cur==-1) {
				ll s=-sum;
				FOR(j,i+1) {
					X[j]-=s;
					sum-=s*A[j];
				}
				out();
				break;
			}
		}
		cur=0;
		for(i=N-1;i>=0;i--) {
			cur+=A[i];
			if(cur==1) {
				ll s=-sum;
				for(j=i;j<N;j++) {
					X[j]+=s;
					sum+=s*A[j];
				}
				out();
				break;
			}
		}
	}
	cout<<"No"<<endl;
	
}

まとめ

これはすんなり思いつけて良かったね。