Codeforces 546E - Soldier and Traveling (最大流)

题意

给出一些城市初始的士兵个数，给出目标的个数，问能不能达到这个目标。

思路

源点到初始建边，容量为初始的个数，目标到汇点建边，容量为目标的个数，他们之间建图，然后跑一下最大流。

数组越界了，报了TLE，吓得我以为dinic写挫了。。

代码

#include <stack>
#include <cstdio>
#include <list>
#include <cassert>
#include <set>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <queue>
#include <functional>
#include <cstring>
#include <algorithm>
#include <cctype>
//#pragma comment(linker, "/STACK:102400000,102400000")
#include <string>
#include <map>
#include <cmath>
//#include <ext/pb_ds/assoc_container.hpp>
//#include <ext/pb_ds/hash_policy.hpp>
using namespace std;
//using namespace __gnu_pbds;
#define LL long long
#define ULL unsigned long long
#define SZ(x) (int)x.size()
#define Lowbit(x) ((x) & (-x))
#define MP(a, b) make_pair(a, b)
#define MS(arr, num) memset(arr, num, sizeof(arr))
#define PB push_back
#define X first
#define Y second
#define ROP freopen("input.txt", "r", stdin);
#define MID(a, b) (a + ((b - a) >> 1))
#define LC rt << 1, l, mid
#define RC rt << 1|1, mid + 1, r
#define LRT rt << 1
#define RRT rt << 1|1
#define FOR(i, a, b) for (int i=(a); (i) < (b); (i)++)
#define FOOR(i, a, b) for (int i = (a); (i)<=(b); (i)++)
const double PI = acos(-1.0);
const int INF = 0x3f3f3f3f;
const double eps = 1e-4;
const int MAXN = 300+10;
const int MOD = 10007;
const int dir[][2] = { {-1, 0}, {1, 0}, {0, -1}, {0, 1} };
const int seed = 131;
int cases = 0;
typedef pair<int, int> pii;

struct EDGE
{
int v, cap, flow, nxt;
};

vector<EDGE> edge;
int st, ed, n, sz;
int vis[MAXN], d[MAXN], cur[MAXN], head[MAXN];
int mp[MAXN][MAXN];
queue<int> Q;

bool bfs()
{
MS(vis, 0);
Q.push(st);
d[st] = 0;
vis[st] = 1;
while (!Q.empty())
{
int u = Q.front(); Q.pop();
for (int i = head[u]; i != -1; i = edge[i].nxt)
{
EDGE &e = edge[i];
if (!vis[e.v] && e.cap > e.flow)
{
vis[e.v] = 1;
d[e.v] = d[u] + 1;
Q.push(e.v);
}
}
}
return vis[ed];
}

int dfs(int u, int a)
{
if (u == ed || a <= 0) return a;
int flow = 0, f;
for (int &i = cur[u]; i != -1; i = edge[i].nxt)
{
EDGE &e = edge[i];
if (d[u] + 1 == d[e.v] && (f = dfs(e.v, min(a, e.cap-e.flow))) > 0)
{
e.flow += f;
edge[i^1].flow -= f;
flow += f;
a -= f;
if (a == 0) break;
}
}
return flow;
}

int sum1;
void Dinic()
{
int flow = 0;
while (bfs())
{
memcpy(cur, head, sizeof head);
flow += dfs(st, INF);
}
if (flow != sum1) { puts("NO"); return; }
puts("YES");
for (int i = 1; i <= n; i++)
{
for (int j = head[i]; j != -1; j = edge[j].nxt)
{
EDGE e = edge[j];
if (e.flow > 0) mp[i][e.v-n] = e.flow;
}
}
for (int i = 1; i <= n; i++)
{
for (int j = 1; j <= n; j++) printf("%d ", mp[i][j]);
puts("");
}
}

void add_edge(int from, int to, int cap)
{
edge.PB({to, cap, 0, head[from]});
edge.PB({from, 0, 0, head[to]});
int sz = SZ(edge);
head[from] = sz-2, head[to] = sz-1;
}

int a[MAXN], b[MAXN];
int main()
{
//ROP;
int m;
scanf("%d%d", &n, &m);
st = 0, ed = n*2+1;
MS(head, -1);
int sum2 = 0;
sum1 = 0;
for (int i = 1; i <= n; i++)
{
scanf("%d", &a[i]);
sum1 += a[i];
add_edge(st, i, a[i]);
add_edge(i, n+i, INF);
}
for (int i = 1; i <= n; i++)
{
scanf("%d", &b[i]);
sum2 += b[i];
add_edge(n+i, ed, b[i]);
}
while (m--)
{
int u, v;
scanf("%d%d", &u, &v);
add_edge(u, v+n, INF); add_edge(v, u+n, INF);
}
if (sum1 != sum2) { puts("NO"); return 0; }
Dinic();
return 0;
}