yukicoder : No.2860 Heal Slimes

コーナーケースを考えるのが難しい。
https://yukicoder.me/problems/no/2860

問題

整数列Hと、整数K,Xが与えられる。
Hの連続するK要素にX加算する、という作業を任意回数行える時、Hの全要素を一致させられるか。

解法

Hの階差数列をAとする。
また、Aの最初と最後に任意の値を設定できる要素を追加できるとする。

Hに対し1回操作を行うということは、Aに対しA[i]にXを加算してA[i+K]からXを減算することに等しい。
A[i]をi%Kで分類してK個の数列に分けたとする。
各数列について以下を満たせるか判定すればよい。

• 各要素はXの倍数
• 総和は0
• prefix sumは常に0以下

Aの最初と最後には任意の値を取れる要素があるので、それらは総和が0となるように調整しよう。

int T,N,K,X;
ll H[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K>>X;
		int mi=1<<30;
		FOR(i,N) {
			cin>>H[i];
		}
		FOR(i,N-1) {
			H[i]=H[i+1]-H[i];
			if(H[i]%X) break;
			H[i]/=X;
		}
		if(i<N-1) {
			cout<<"No"<<endl;
			continue;
		}
		FOR(i,K) {
			if(i==N-1) continue;
			ll sum=0;
			for(j=i;j<N-1;j+=K) sum+=H[j];
			if(i==K-1) {
				if((N-1)%K==i) continue;
				if(sum>=0) H[i]-=sum;
			}
			else if((N-1)%K==i) {
				if(sum<=0) H[j-K]-=sum;
			}
			sum=0;
			for(j=i;j<N-1;j+=K) {
				sum+=H[j];
				if(sum>0) break;
			}
			if(sum!=0) break;
		}
		if(i==K) {
			cout<<"Yes"<<endl;
		}
		else {
			cout<<"No"<<endl;
		}
		
	}
}

まとめ

今回なんで妙にスライム推しなんだろうな。