AtCoder ABC #307 (東京海上日動プログラミングコンテスト2023) : G - Approximate Equalization

割とてこずった。
https://atcoder.jp/contests/abc307/tasks/abc307_g

問題

N要素の整数列Aが与えられる。
連続する2要素のうち、片方をインクリメント、片方をデクリメントする、という処理を繰り返す。
Aの最大値と最小値の差が1以下となるようにするには、最小何回処理を行う必要があるか。

解法

まずAの各要素から平均値を引いておく。
そうすると、最終的にAは0/1のいずれかだけ取るようにする問題と置くことができる。

AのprefixSumを考えると、問題の処理は末尾以外の1要素をインクリメントまたはデクリメントすることに相当する。

最終的に、A中1の数が残す数をXとする。
dp(n,k) := Aのprefix n要素まで、結果の0/1を確定させ、計k個1があるとき、それを達成する最小処理回数
とすると、このテーブルをO(N^2)で埋めて、dp(N,X)を答えればよい。

int N;
ll A[5050];

ll from[5050],to[5050];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	ll S=0;
	FOR(i,N) {
		cin>>A[i];
		A[i]+=1LL<<30;
		S+=A[i];
	}
	FOR(i,N) {
		A[i]-=S/N;
		if(i) A[i]+=A[i-1];
	}
	S%=N;
	from[0]=0;
	FOR(i,N+1) from[i]=1LL<<60;
	from[0]=0;
	FOR(i,N-1) {
		FOR(j,N+1) to[j]=1LL<<60;
		FOR(j,N+1) {
			to[j]=min(to[j],from[j]+abs(A[i]-j));
			to[j+1]=min(to[j+1],from[j]+abs(A[i]-j-1));
		}
		swap(from,to);
	}
	
	if(S) {
		cout<<min(from[S],from[S-1])<<endl;
	}
	else {
		cout<<from[S]<<endl;
	}
	
}

まとめ

ちょっと自信なかったけど、どうにか解けてよかったね。