我TLE代码：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
ll n,m,s,t,idx;
const ll N=1001*2+10,M=2500*2+10;
ll maxflow;
struct edge{
	ll nxt,to,dis,c;
} e[M];
ll head[M];
void add(long long u,long long v,long long fl,long long d) {
	e[++idx].to=v;
	e[idx].nxt=head[u];
	e[idx].c=fl;
	e[idx].dis=d;
	head[u]=idx;
}
ll vis[N],pre[N],dist[N],ans2;
ll incf[N],sum,a[N];
const ll inf=2e17;
bool bfs() {
	memset(vis,0,sizeof(vis));
	queue <ll> q;
	q.push(s);
	memset(dist,0x3f,sizeof(dist));
	vis[s]=1;
	incf[s]=inf;
	dist[s]=0;
	while(q.size()) {
		ll u=q.front();
		q.pop();
		vis[u]=0;
		for(register ll i=head[u];i;i=e[i].nxt) {
			ll v=e[i].to;
			if(e[i].c&&dist[v]>dist[u]+e[i].dis) {
				dist[v]=dist[u]+e[i].dis;
				incf[v]=min(incf[u],e[i].c);
				pre[v]=i;
				if(vis[v])
				continue;
				vis[v]=1;
				q.push(v);
			}
		}
	}
	return dist[t]!=dist[0];
}
void update() {
	ll u=t;
	while(u!=s) {
		ll i=pre[u];
		e[i].c-=incf[t];// i正向 i^1反向 
		e[i^1].c+=incf[t];
		u=e[i^1].to;
	}
	maxflow+=incf[t];
	ans2+=incf[t]*dist[t];
}
signed main() {
	scanf("%lld",&n);
	s=1,t=n+5;
	for(register int i=2;i<=n+1;i++) {
		scanf("%lld",&a[i]);
		sum+=a[i];
	}
	sum/=n;
	for(register int i=2;i<=n+1;i++) {
		if(a[i]>sum) {
			add(s,i,a[i]-sum,0);
			add(i,s,0,0);
		}
		else if(a[i]<sum) {
			add(i,t,sum-a[i],0);
			add(t,i,0,0);
		}
	}
	for(int i=2;i<=n+1;i++) {
		if(i!=n+1)
		add(i,i+1,inf,1),add(i+1,i,0,-1);
		if(i!=2)
		add(i,i-1,inf,1),add(i-1,i,0,-1);
	}
	add(n+1,2,inf,1),add(2,n+1,0,-1);
	add(2,n+1,inf,1),add(n+1,2,0,-1);
	while(bfs()) {
		update();
	}
	printf("%lld\n",ans2);
	return 0;
}