【CF739E】Gosha is hunting

题目

题目链接：https://codeforces.com/problemset/problem/739/E
你要抓神奇宝贝！
现在一共有 $N$ 只神奇宝贝。
你有 $a$ 个『宝贝球』和 $b$ 个『超级球』。
『宝贝球』抓到第 $i$ 只神奇宝贝的概率是 $p_i$ ，『超级球』抓到的概率则是 $u_i$ 。
不能往同一只神奇宝贝上使用超过一个同种的『球』，但是可以往同一只上既使用『宝贝球』又使用『超级球』（都抓到算一个）。
请合理分配每个球抓谁，使得你抓到神奇宝贝的总个数期望最大，并输出这个值。
$n\le 2000$。

思路

考虑建立费用流模型。首先对于宝贝球和超级球分别建立一个点，从源点连一条流量为球的数量，费用为 $0$ 的边来限制每一个球的数量。
然后对于一个点 $i$，从宝贝球的点连一条 $(1,p_i)$ 的边，从超级球连一条 $(1,u_i)$ 的边，然后从 $i$ 向汇点连一条 $(1,0)$ 的边，再连一条 $(1,-p_i\times u_i)$ 的边。然后跑最大费用最大流。如果这个逼只用了一个球来抓，那么就会优先流费用为 $0$ 的边，否则两条边都流，贡献即为 $p_i+u_i-p_i{u}_i$，恰好是抓到它的期望。
时间复杂度 $O(n^3)$。

代码

#include <bits/stdc++.h>
using namespace std;

const int N=2010,M=20010,Inf=1e9;
const double eps=1e-10;
int n,a,b,A,B,S,T,tot,pre[N],head[N];
double cost,p[2][N],dis[N];
bool vis[N];

struct edge
{
	int next,to,flow;
	double cost;
}e[M];

void add(int from,int to,int flow,double cost)
{
	e[++tot]=(edge){head[from],to,flow,cost};
	head[from]=tot;
	swap(from,to);
	e[++tot]=(edge){head[from],to,0,-cost};
	head[from]=tot;
}

bool spfa()
{
	memset(dis,0xcf,sizeof(dis));
	queue<int> q;
	q.push(S); dis[S]=0;
	while (q.size())
	{
		int u=q.front(); q.pop();
		vis[u]=0;
		for (int i=head[u];~i;i=e[i].next)
		{
			int v=e[i].to;
			if (e[i].flow && dis[v]<dis[u]+e[i].cost-eps)
			{
				dis[v]=dis[u]+e[i].cost; pre[v]=i;
				if (!vis[v]) { vis[v]=1; q.push(v); }
			}
		}
	}
	return dis[T]>0;
}

void addflow()
{
	int minf=Inf;
	for (int i=T;i!=S;i=e[pre[i]^1].to)
		minf=min(minf,e[pre[i]].flow);
	for (int i=T;i!=S;i=e[pre[i]^1].to)
	{
		e[pre[i]].flow-=minf;
		e[pre[i]^1].flow+=minf;
	}
	cost+=dis[T]*minf;
}

void MCMF()
{
	while (spfa()) addflow();
}

int main()
{
	memset(head,-1,sizeof(head));
	tot=1; S=N-1; T=N-2; A=N-3; B=N-4;
	scanf("%d%d%d",&n,&a,&b);
	for (int i=1;i<=n;i++) scanf("%lf",&p[0][i]);
	for (int i=1;i<=n;i++) scanf("%lf",&p[1][i]);
	for (int i=1;i<=n;i++)
	{
		add(A,i,1,p[0][i]); add(B,i,1,p[1][i]);
		add(i,T,1,0); add(i,T,1,-p[0][i]*p[1][i]);
	}
	add(S,A,a,0); add(S,B,b,0);
	MCMF();
	printf("%.10lf",cost);
	return 0;
}