负载平衡问题

https://www.acwing.com/problem/content/2196/

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
const int N=5010,M=200010,inf=1e8;
int S,T;
int h[N],e[M],f[M],w[M],ne[M],idx;
int q[N],d[N],pre[N],incf[N];
bool st[N];
int a[110];
void add(int a,int b,int c,int d)
{
	e[idx]=b;f[idx]=c;w[idx]=d;ne[idx]=h[a];h[a]=idx++;
	e[idx]=a;f[idx]=0;w[idx]=-d;ne[idx]=h[b];h[b]=idx++;
}
bool spfa()
{
	int hh=0,tt=1;
	memset(d,0x3f,sizeof d);
	memset(incf,0,sizeof incf);
	q[0]=S;d[S]=0,incf[S]=inf;
	while(hh!=tt)
	{
		int t=q[hh++];
		if(hh==N) hh=0;
		st[t]=0;
		for(int i=h[t];i!=-1;i=ne[i])
		{
			int ver=e[i];
			if(f[i]&&d[ver]>d[t]+w[i])
			{
				d[ver]=d[t]+w[i];
				pre[ver]=i;
				incf[ver]=min(f[i],incf[t]);
				if(!st[ver])
				{
					q[tt++]=ver;
					if(tt==N) tt=0;
					st[ver]=1;
				}
			}
		}
	}
	return incf[T]>0;
}
void EK(int &flow,int &cost)
{
	flow=cost=0;
	while(spfa())
	{
		int t=incf[T];
		flow+=t,cost+=t*d[T];
		for(int i=T;i!=S;i=e[pre[i]^1])
		{
			f[pre[i]]-=t;
			f[pre[i]^1]+=t;
		}
	}
}
int main()
{
	int n,m;
	scanf("%d",&n);
	memset(h,-1,sizeof h);
	S=0,T=n+1;
	int sum=0;
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&a[i]);
		sum+=a[i];
		
	} 
	sum/=n;
	for(int i=1;i<=n;i++)
	{
		if(a[i]>sum)
		{
			add(S,i,a[i]-sum,0);
		}
		else if(sum>a[i])
		{
			add(i,T,sum-a[i],0);
		}
		
	}
	for(int i=2;i<=n-1;i++)
	{
		add(i,i+1,inf,1);
		add(i,i-1,inf,1);
	}
	add(1,2,inf,1);
	add(1,n,inf,1);
	add(n,1,inf,1);
	add(n,n-1,inf,1);
	int flow,cost;
	EK(flow,cost);
	printf("%d\n",cost);

	return 0;
}