kmjp's blog

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

lay_contest beta round 00004 : C - 神による整地 (3)

意外にすんなり解けた。
https://www.hackerrank.com/contests/lay-contest-00004/challenges/3-3

問題

R*Cのグリッドに対し、各セル値H[r][c]が与えられる。
合計最大N回まで、どこかのセルの値を1増加してどこかのセルの値を1減少できる。
H[r][c]の2乗和を最小化せよ。

解法

全体を平均化するほど二乗和が最小になる。
各セルの差がそれ以上平らにならない、すなわち最大値と最小値の差が1以下にするのにN回以内の増減で達成可能なら、全セルを平均値にしよう。
(最大値と最小値に1の差が生じる場合、元々Hの小さいセルに小さい値を割り当てればよい)

N回以内の増減で達成できない場合、値を大きい順にN個減少させ、小さい順にN個増加させる。
もちろんN回愚直に増減をシミュレートすると間に合わないから、最大値が同じ値がp個ある場合、p回の処理で全要素を1減らす、ということを上から(p+1)個目の要素の値になるまで繰り返すことを行えばよい。

int T;
int M;
ll N;
ll A[505*505];
ll U[505*505];

void solve() {
	int i,j,k,l,r,x,y; string s;
	cin>>T;
	while(T--) {
		cin>>x>>y>>N;
		M=x*y;
		ll S=0;
		FOR(i,M) cin>>A[i], S+=A[i];
		if(S<0) {
			FOR(i,M) A[i]=-A[i];
			S=-S;
		}
		sort(A,A+M);
		ll num=0;
		FOR(i,M) {
			U[i]=S/M+(i<(S%M));
			if(A[i]>U[i]) num+=A[i]-U[i];
		}
		
		if(num<=N) FOR(i,M) A[i]=U[i];
		else {
			FOR(x,2) {
				int L=M-1;
				ll S=N;
				while(1) {
					if((A[L]-A[L-1])*(M-L)<=S) {
						S-=(A[L]-A[L-1])*(M-L);
						L--;
					}
					else {
						for(i=M-L-1;i>=0;i--) A[L+i]=A[L]-(S/(M-L)+(i<S%(M-L)));
						break;
					}
				}
				reverse(A,A+M);
				FOR(i,M) A[i]=-A[i];
			}
		}
		
		
		ll tot=0;
		FOR(i,M) tot+=1LL*A[i]*A[i];
		cout<<tot<<endl;
		
	}

}

まとめ

神と整地と言ってもアクトレイザーではないよなぁ。