G - Approximate Equalization Editorial by purslane

A much ( ? ) Easier Solution

First of all , we can find that if we perform an operation , \(\sum_{i=1}^n a_i\) won’t change. So we know the final \(\sum_{i=1}^n a'_i=\sum_{i=1}^n a_i\) . What’s more , when \(\forall 1 \le i,j \le n , |a'_i-a'_j| \le 1\) holds , \(a'_i\) can only be two situations : \(\lfloor \frac{S}{n} \rfloor\) and \(\lfloor \frac{S}{n} \rfloor+1\) . And the number of the second situation is exactly \(S-n \times \lfloor \frac{S}{n} \rfloor\) .

We can try to sort all the operations by the position it performs on . Then we can use a simple dp .

Assume \(dp_{x,y}\) represents the minimum operations needed to make \(y\) of the previous \(x\) elements of the array are \(\lfloor \frac{S}{n} \rfloor+1\) and the others of the previous \(x\) elements are \(\lfloor \frac{S}{n} \rfloor\) .

When we do transfers, we are supposed to know when the previous \(x\) elements have \(y\) second situations , what \(a_{x+1}\) will be . Well , now we have only performed the operations which can only influence the previous \(x+1\) elements , so the sum of them won’t change . Assume \(pre_x\) to be \(\sum_{i=1}^x a_i\) . So we know that \(a'_{x+1} = pre_{x+1} - y \times ( \lfloor \frac{S}{n} \rfloor+1) - (x-y) \times \lfloor \frac{S}{n} \rfloor\) . So \(dp_{x+1,y}\) can come from \(dp_{x,y}+|a'_{x+1}-\lfloor \frac{S}{n} \rfloor|\) , \(dp_{x+1,y+1}\) can come from \(dp_{x,y}+|a'_{x+1}-\lfloor \frac{S}{n} \rfloor - 1|\) .

Therefore , the problem is solved within \(O(n^2)\) time , which is fast sufficiently.

Here is my code for you to refer to.

#include<bits/stdc++.h>
#define int long long
#define ffor(i,a,b) for(int i=(a);i<=(b);i++)
#define roff(i,a,b) for(int i=(a);i>=(b);i--)
using namespace std;
const int MAXN=5000+10;
int n,v,a[MAXN],pre[MAXN],tot,dp[MAXN][MAXN];
signed main() {
	ios::sync_with_stdio(false),cin.tie(0),cout.tie(0);
	cin>>n; ffor(i,1,n) cin>>a[i],tot+=a[i],pre[i]=pre[i-1]+a[i];
	int v=tot/n; if(v*n>tot) v--;
	int x=n-(tot-v*n); 
	memset(dp,0x3f,sizeof(dp));
	dp[0][0]=0;
	ffor(i,0,n-1) ffor(j,0,i) {
		int val=j*(v+1)+(i-j)*v,minus=val-pre[i],A=a[i+1]-minus;
		dp[i+1][j]=min(dp[i+1][j],dp[i][j]+abs(A-v)),dp[i+1][j+1]=min(dp[i+1][j+1],dp[i][j]+abs(A-(v+1)));
	}
	cout<<dp[n][n-x];
	return 0;
}

posted:
last update: