Codeforces #1019 : Div2 E. Keep the Sum

問題文は短いけど、コード量は結構増えてしまった。
https://codeforces.com/contest/2103/problem/E

問題

正整数Kと、N要素の整数列Aが与えられる。Aの値は、0以上K以下である。
以下を繰り返し、Aを3N手以下で単調増加にできるか、できるなら一例を示せ。

• A[i]+A[j]=Kとなる2要素(i,j)を選び、和がKでかつどちらも0以上K以下となる範囲で、両者の値を任意に変更する。

解法

初期状態ですでに単調増加なら何もしない。
そうでないとき、初手で何もできないなら解なし。

初手で操作が可能なら、0～Kの値を何でも作れることになる。
あとは、後ろにできるだけK、前に0を詰めるようにしていく。

int T;
ll N,K;
ll A[202020];


vector<array<int,3>> ret;
void go(int a,int b,int v) {
	if(v==0) return;
	assert(A[a]+A[b]==K);
	assert(A[a]>=0&&A[b]>=0);
	assert(v<=A[a]&&v>=-A[b]);
	assert(a!=b);
	ret.push_back({a+1,b+1,v});
	A[a]-=v;
	A[b]+=v;
	assert(A[a]>=0&&A[b]>=0);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		FOR(i,N) cin>>A[i];
		FOR(i,N-1) if(A[i]>A[i+1]) break;
		if(i==N-1) {
			cout<<0<<endl;
			continue;
		}
		ret.clear();
		map<int,int> M;
		x=y=-1;
		int zero=-1,full=-1;
		FOR(i,N) {
			if(M.count(K-A[i])) {
				zero=M[K-A[i]];
				full=i;
			}
			M[A[i]]=i;
		}
		if(zero==-1) {
			cout<<-1<<endl;
			continue;
		}
		if(full!=N-1) {
			go(zero,full,A[zero]-(K-A[N-1]));
			full=N-1;
		}
		set<pair<int,int>> S;
		FOR(i,N) {
			if(i!=zero&&i!=full) S.insert({A[i],i});
		}
		for(i=N-2;i>=2;i--) {
			x=S.rbegin()->second;
			S.erase({A[x],x});
			if(x==i) continue;
			if(zero==i) {
				go(i,full,A[i]-(K-A[x]));
				go(x,i,A[x]-A[i]);
				swap(zero,x);
			}
			else if(zero!=i) {
				S.erase({A[i],i});
				go(zero,full,A[zero]-A[i]);
				go(i,full,A[i]-(K-A[x]));
				go(x,i,A[x]-A[i]);
				swap(zero,x);
				S.insert({A[x],x});
			}
		}
		
		if(zero==0) {
			go(zero,full,A[zero]);
		}
		else {
			go(1,full,A[1]-A[0]);
			go(0,full,A[0]);
		}
		
		FOR(i,N-1) assert(A[i]<=A[i+1]);
		cout<<ret.size()<<endl;
		FORR(a,ret) {
			cout<<a[0]<<" "<<a[1]<<" "<<a[2]<<endl;
		}
		
		
		
	}
}

まとめ

考え方はシンプルだけど、実装はちょっとめんどい。