CF Round #304 Div.2 E. Soldier and Traveling

Codeforces Round #304 Div.2 E. Soldier and Traveling

Problem - 546E - Codeforces

フローで解けることを知っていながら解いたが、解説無しでACしたのでよし。

n個の街があって、m個の道によって連結である。各街には兵士がa[i]人かいて、彼らは今いる街にとどまってもよいし、1個だけ道を使って別の街に移動しても良い。このとき、各街について、定められた兵士の数b[i]人にすることができるかを答え、できるならその移動の仕方を答える、という問題。

まず、「ある点Pに兵士がx人いて、それをy人に変える」という状況を考えるために、点Sと点Tを別で用意し、SからPに容量x, PからTに容量yの辺を張る、ということをとりあえず考えてみた。そして、点Pと点Qが繋がっているという状況をどう作るかを考えてみると、単純に容量a[P]、a[Q]でつなぐ、という解法を思いついたが、それでは兵士が何回も移動できることになってしまうのでダメだった。そこで、Pから移動できる人数は合計でa[P]である、ということに注目すると、まずPの中継点P'に容量a[P]の辺を張って、そこから各Qに対して容量INFの辺を張ることで、その点から流れるフローの大きさの合計がもともといた人数以下におさえることができる。そしてうまくいった。

あとで他の人の解説を見ると、辺の張り方が少し違っていたが、本質的には同じことをしていた。

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

struct edge {
        int to, cap, rev;
};
bool used[101010];
vector<edge> g[101010];
static const int INF = 0x3f3f3f3f;
void add_edge(int from, int to, int cap) {
        g[from].push_back((edge) { to, cap, (int)g[to].size() });
        g[to].push_back((edge) { from, 0, (int)g[from].size() - 1 });
}
int dfs(int v, int t, int f) {
        if (v == t) return f;
        used[v] = true;
        for (int i = 0; i < g[v].size(); i ++) {
                edge& e = g[v][i];
                if (!used[e.to] && e.cap > 0) {
                        int d = dfs(e.to, t, min(f, e.cap));
                        if (d > 0) {
                                e.cap -= d;
                                g[e.to][e.rev].cap += d;
                                return  d;
                        }
                }
        }
        return 0;
}
int MaxFlow(int s, int t) {
        int flow = 0;
        while (true) {
                memset(used, false, sizeof(used));
                int f = dfs(s, t, INF);
                if (f == 0) return flow;
                flow += f;
        }
}

int main() {
        int n, m;
        cin >> n >> m;
        vector<int> a(n), b(n);
        int sum1 = 0, sum2 = 0;
        for (int i = 0; i < n; i ++) { 
                cin >> a[i];
                sum1 += a[i];
        }
        for (int i = 0; i < n; i ++) { 
                cin >> b[i];
                sum2 += b[i];
        }
        if (sum1 != sum2) { //総人数は増減しないので異なるとダメ
                cout << "NO" << endl;
                return 0;
        }
        int s = 2 * n, t = 2 * n + 1;
        for (int i = 0; i < n; i ++) {
                add_edge(s, i, a[i]);
                add_edge(i, t, b[i]);
                add_edge(i, i + n, a[i]);  //中継点にa[i]を張っておく
        }
        for (int i = 0; i < m; i ++) {
                int x, y;
                cin >> x >> y;
                x --, y --;
                add_edge(x + n, y, INF); //中継点から目的の点に張る
                add_edge(y + n, x, INF);
        }
        if (MaxFlow(s, t) != sum1) { //全部流せていないとダメ
                cout << "NO" << endl;
                return 0;
        } 
        cout << "YES" << endl;
        vector<vector<int>> ans(n, vector<int> (n));
        for (int i = 0; i < n; i ++) {
                for (int j = 0; j < n; j ++) {
                        if (i == j) {
                                for (auto e : g[i]) {
                                        if (e.to == i + n) {
                                                ans[i][j] = e.cap; //中継点に向かって流さなかった分（まだ流せる分）がそこにとどまる人数
                                        }
                                }
                        } else {
                                bool found = false;
                                for (auto e : g[i + n]) {
                                        if (e.to >= n) continue;
                                        if (e.to == j) {
                                                found = true;
                                                ans[i][j] = g[e.to][e.rev].cap; //流した分は、逆辺の容量と等しい
                                                break;
                                        }
                                }
                                if (!found) ans[i][j] = 0;
                        }
                }
        }
        for (int i = 0; i < n; i ++) {
                for (int j = 0; j < n; j ++) {
                        cout << ans[i][j] << (j == n - 1 ? '\n' : ' ');
                }
        }
        return 0;
}