mcmf.hpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | #include<bits/stdc++.h> | ||
| 3 | // #include<bits/extc++.h> | ||
| 4 | #include <ext/pb_ds/priority_queue.hpp> | ||
| 5 | |||
| 6 | // NOTE: This doesn't support negative-cost edges; you can adjust edge weights | ||
| 7 | // (e.g. by precomputing a potential function) to make them positive. | ||
| 8 | |||
| 9 | template <typename flow_t = int, typename cost_t = int64_t> | ||
| 10 | struct MCMF_SSPA { | ||
| 11 | int N; | ||
| 12 | std::vector<std::vector<int>> adj; | ||
| 13 | struct edge_t { | ||
| 14 | int dest; | ||
| 15 | flow_t cap; | ||
| 16 | cost_t cost; | ||
| 17 | }; | ||
| 18 | std::vector<edge_t> edges; | ||
| 19 | |||
| 20 | std::vector<char> seen; | ||
| 21 | std::vector<cost_t> pi; | ||
| 22 | std::vector<int> prv; | ||
| 23 | |||
| 24 | ✗ | explicit MCMF_SSPA(int N_) : N(N_), adj(N), pi(N, 0), prv(N) {} | |
| 25 | |||
| 26 | ✗ | void add_edge(int from, int to, flow_t cap, cost_t cost) { | |
| 27 | ✗ | assert(cap >= 0); | |
| 28 | ✗ | assert(cost + pi[from] - pi[to] >= 0); // TODO: Remove this restriction | |
| 29 | ✗ | int e = int(edges.size()); | |
| 30 | ✗ | edges.emplace_back(edge_t{to, cap, cost}); | |
| 31 | ✗ | edges.emplace_back(edge_t{from, 0, -cost}); | |
| 32 | ✗ | adj[from].push_back(e); | |
| 33 | ✗ | adj[to].push_back(e+1); | |
| 34 | } | ||
| 35 | |||
| 36 | static constexpr cost_t INF_COST = std::numeric_limits<cost_t>::max() / 4; | ||
| 37 | static constexpr flow_t INF_FLOW = std::numeric_limits<flow_t>::max() / 4; | ||
| 38 | std::vector<cost_t> dist; | ||
| 39 | __gnu_pbds::priority_queue<std::pair<cost_t, int>> q; | ||
| 40 | std::vector<typename decltype(q)::point_iterator> its; | ||
| 41 | ✗ | cost_t dijkstra(int s, int t) { | |
| 42 | ✗ | dist.assign(N, INF_COST); | |
| 43 | ✗ | dist[s] = 0; | |
| 44 | |||
| 45 | ✗ | its.assign(N, q.end()); | |
| 46 | ✗ | its[s] = q.push({-(dist[s] - pi[s]), s}); | |
| 47 | |||
| 48 | ✗ | while (!q.empty()) { | |
| 49 | ✗ | int i = q.top().second; q.pop(); | |
| 50 | ✗ | cost_t d = dist[i]; | |
| 51 | ✗ | for (int e : adj[i]) { | |
| 52 | ✗ | if (edges[e].cap) { | |
| 53 | ✗ | int j = edges[e].dest; | |
| 54 | ✗ | cost_t nd = d + edges[e].cost; | |
| 55 | ✗ | if (nd < dist[j]) { | |
| 56 | ✗ | dist[j] = nd; | |
| 57 | ✗ | prv[j] = e; | |
| 58 | ✗ | if (its[j] == q.end()) { | |
| 59 | ✗ | its[j] = q.push({-(dist[j] - pi[j]), j}); | |
| 60 | } else { | ||
| 61 | ✗ | q.modify(its[j], {-(dist[j] - pi[j]), j}); | |
| 62 | } | ||
| 63 | } | ||
| 64 | } | ||
| 65 | } | ||
| 66 | } | ||
| 67 | |||
| 68 | ✗ | swap(pi, dist); | |
| 69 | ✗ | return pi[t]; | |
| 70 | } | ||
| 71 | |||
| 72 | ✗ | flow_t path(int s, int t) { | |
| 73 | ✗ | flow_t cur_flow = std::numeric_limits<flow_t>::max(); | |
| 74 | ✗ | for (int cur = t; cur != s; ) { | |
| 75 | ✗ | int e = prv[cur]; | |
| 76 | ✗ | int nxt = edges[e^1].dest; | |
| 77 | ✗ | cur_flow = std::min(cur_flow, edges[e].cap); | |
| 78 | ✗ | cur = nxt; | |
| 79 | } | ||
| 80 | ✗ | for (int cur = t; cur != s; ) { | |
| 81 | ✗ | int e = prv[cur]; | |
| 82 | ✗ | int nxt = edges[e^1].dest; | |
| 83 | ✗ | edges[e].cap -= cur_flow; | |
| 84 | ✗ | edges[e^1].cap += cur_flow; | |
| 85 | ✗ | cur = nxt; | |
| 86 | } | ||
| 87 | ✗ | return cur_flow; | |
| 88 | } | ||
| 89 | |||
| 90 | ✗ | std::vector<std::pair<flow_t, cost_t>> all_flows(int s, int t, cost_t max_cost = INF_COST - 1) { | |
| 91 | ✗ | assert(s != t); | |
| 92 | ✗ | std::vector<std::pair<flow_t, cost_t>> res; | |
| 93 | ✗ | while (dijkstra(s, t) <= max_cost) { | |
| 94 | ✗ | assert(res.empty() || pi[t] >= res.back().second); | |
| 95 | ✗ | flow_t f = path(s, t); | |
| 96 | ✗ | res.push_back({f, pi[t]}); | |
| 97 | } | ||
| 98 | ✗ | return res; | |
| 99 | } | ||
| 100 | |||
| 101 | ✗ | std::pair<flow_t, cost_t> max_flow(int s, int t, cost_t max_cost = INF_COST - 1) { | |
| 102 | ✗ | assert(s != t); | |
| 103 | ✗ | flow_t tot_flow = 0; cost_t tot_cost = 0; | |
| 104 | ✗ | while (dijkstra(s, t) <= max_cost) { | |
| 105 | ✗ | flow_t cur_flow = path(s, t); | |
| 106 | ✗ | tot_flow += cur_flow; | |
| 107 | ✗ | tot_cost += cur_flow * pi[t]; | |
| 108 | } | ||
| 109 | ✗ | return {tot_flow, tot_cost}; | |
| 110 | } | ||
| 111 | }; | ||
| 112 | |||
| 113 | template <typename flow_t = int, typename cost_t = int64_t> | ||
| 114 | struct MCMF_Dinic { | ||
| 115 | int N; | ||
| 116 | std::vector<std::vector<int>> adj; | ||
| 117 | struct edge_t { | ||
| 118 | int dest; | ||
| 119 | flow_t cap; | ||
| 120 | cost_t cost; | ||
| 121 | }; | ||
| 122 | std::vector<edge_t> edges; | ||
| 123 | |||
| 124 | std::vector<char> seen; | ||
| 125 | std::vector<cost_t> pi; | ||
| 126 | |||
| 127 | ✗ | explicit MCMF_Dinic(int N_) : N(N_), adj(N), pi(N, 0) {} | |
| 128 | |||
| 129 | ✗ | void add_edge(int from, int to, flow_t cap, cost_t cost) { | |
| 130 | ✗ | assert(cap >= 0); | |
| 131 | ✗ | assert(cost + pi[from] - pi[to] >= 0); // TODO: Remove this restriction | |
| 132 | ✗ | int e = int(edges.size()); | |
| 133 | ✗ | edges.emplace_back(edge_t{to, cap, cost}); | |
| 134 | ✗ | edges.emplace_back(edge_t{from, 0, -cost}); | |
| 135 | ✗ | adj[from].push_back(e); | |
| 136 | ✗ | adj[to].push_back(e+1); | |
| 137 | } | ||
| 138 | |||
| 139 | static constexpr cost_t INF_COST = std::numeric_limits<cost_t>::max() / 4; | ||
| 140 | static constexpr flow_t INF_FLOW = std::numeric_limits<flow_t>::max() / 4; | ||
| 141 | std::vector<cost_t> dist; | ||
| 142 | __gnu_pbds::priority_queue<std::pair<cost_t, int>> q; | ||
| 143 | std::vector<typename decltype(q)::point_iterator> its; | ||
| 144 | ✗ | cost_t dijkstra(int s, int t) { | |
| 145 | ✗ | dist.assign(N, INF_COST); | |
| 146 | ✗ | dist[s] = 0; | |
| 147 | |||
| 148 | ✗ | its.assign(N, q.end()); | |
| 149 | ✗ | its[s] = q.push({-(dist[s] - pi[s]), s}); | |
| 150 | |||
| 151 | ✗ | while (!q.empty()) { | |
| 152 | ✗ | int i = q.top().second; q.pop(); | |
| 153 | ✗ | cost_t d = dist[i]; | |
| 154 | ✗ | for (int e : adj[i]) { | |
| 155 | ✗ | if (edges[e].cap) { | |
| 156 | ✗ | int j = edges[e].dest; | |
| 157 | ✗ | cost_t nd = d + edges[e].cost; | |
| 158 | ✗ | if (nd < dist[j]) { | |
| 159 | ✗ | dist[j] = nd; | |
| 160 | ✗ | if (its[j] == q.end()) { | |
| 161 | ✗ | its[j] = q.push({-(dist[j] - pi[j]), j}); | |
| 162 | } else { | ||
| 163 | ✗ | q.modify(its[j], {-(dist[j] - pi[j]), j}); | |
| 164 | } | ||
| 165 | } | ||
| 166 | } | ||
| 167 | } | ||
| 168 | } | ||
| 169 | |||
| 170 | ✗ | std::swap(pi, dist); | |
| 171 | ✗ | return pi[t]; | |
| 172 | } | ||
| 173 | |||
| 174 | std::vector<int> buf; | ||
| 175 | std::vector<int> level; | ||
| 176 | ✗ | flow_t dinic_dfs(int cur, int t, flow_t f) { | |
| 177 | ✗ | if (cur == t) return f; | |
| 178 | ✗ | flow_t cur_f = 0; | |
| 179 | ✗ | assert(f > 0); | |
| 180 | ✗ | for (; buf[cur] < int(adj[cur].size()); buf[cur]++) { | |
| 181 | ✗ | int e = adj[cur][buf[cur]]; | |
| 182 | ✗ | int nxt = edges[e].dest; | |
| 183 | ✗ | if (level[nxt] == level[cur] + 1 && edges[e].cap > 0 && edges[e].cost == pi[nxt] - pi[cur]) { | |
| 184 | ✗ | flow_t v = dinic_dfs(nxt, t, std::min(f, edges[e].cap)); | |
| 185 | ✗ | edges[e].cap -= v; | |
| 186 | ✗ | edges[e^1].cap += v; | |
| 187 | ✗ | f -= v; | |
| 188 | ✗ | cur_f += v; | |
| 189 | ✗ | if (f == 0) break; | |
| 190 | } | ||
| 191 | } | ||
| 192 | ✗ | return cur_f; | |
| 193 | } | ||
| 194 | ✗ | flow_t dinic(int s, int t) { | |
| 195 | ✗ | flow_t tot_flow = 0; | |
| 196 | ✗ | while (true) { | |
| 197 | ✗ | buf.clear(); | |
| 198 | ✗ | buf.reserve(N); | |
| 199 | ✗ | level.assign(N, -1); | |
| 200 | ✗ | buf.push_back(s); | |
| 201 | ✗ | level[s] = 0; | |
| 202 | ✗ | for (int z = 0; z < int(buf.size()); z++) { | |
| 203 | ✗ | int cur = buf[z]; | |
| 204 | ✗ | for (int e : adj[cur]) { | |
| 205 | ✗ | int nxt = edges[e].dest; | |
| 206 | ✗ | if (edges[e].cap > 0 && edges[e].cost == pi[nxt] - pi[cur] && level[nxt] == -1) { | |
| 207 | ✗ | level[nxt] = level[cur] + 1; | |
| 208 | ✗ | buf.push_back(nxt); | |
| 209 | } | ||
| 210 | } | ||
| 211 | } | ||
| 212 | ✗ | if (level[t] == -1) break; | |
| 213 | ✗ | buf.assign(N, 0); | |
| 214 | ✗ | tot_flow += dinic_dfs(s, t, INF_FLOW); | |
| 215 | } | ||
| 216 | ✗ | return tot_flow; | |
| 217 | } | ||
| 218 | |||
| 219 | ✗ | std::vector<std::pair<flow_t, cost_t>> all_flows(int s, int t, cost_t max_cost = INF_COST - 1) { | |
| 220 | ✗ | assert(s != t); | |
| 221 | ✗ | std::vector<std::pair<flow_t, cost_t>> res; | |
| 222 | ✗ | while (dijkstra(s, t) <= max_cost) { | |
| 223 | ✗ | assert(res.empty() || pi[t] > res.back().second); | |
| 224 | ✗ | flow_t f = dinic(s, t); | |
| 225 | ✗ | res.push_back({f, pi[t]}); | |
| 226 | } | ||
| 227 | ✗ | return res; | |
| 228 | } | ||
| 229 | |||
| 230 | ✗ | std::pair<flow_t, cost_t> max_flow(int s, int t, cost_t max_cost = INF_COST - 1) { | |
| 231 | ✗ | assert(s != t); | |
| 232 | ✗ | flow_t tot_flow = 0; cost_t tot_cost = 0; | |
| 233 | ✗ | while (dijkstra(s, t) <= max_cost) { | |
| 234 | ✗ | flow_t cur_flow = dinic(s, t); | |
| 235 | ✗ | tot_flow += cur_flow; | |
| 236 | ✗ | tot_cost += cur_flow * pi[t]; | |
| 237 | } | ||
| 238 | ✗ | return {tot_flow, tot_cost}; | |
| 239 | } | ||
| 240 | }; | ||
| 241 | |||
| 242 | template <typename flow_t = int, typename tot_flow_t = flow_t> | ||
| 243 | struct Dinic { | ||
| 244 | int N; | ||
| 245 | std::vector<std::vector<int>> adj; | ||
| 246 | struct edge_t { | ||
| 247 | int dest; | ||
| 248 | flow_t cap; | ||
| 249 | }; | ||
| 250 | std::vector<edge_t> edges; | ||
| 251 | |||
| 252 | std::vector<char> seen; | ||
| 253 | |||
| 254 | ✗ | explicit Dinic(int N_) : N(N_), adj(N) {} | |
| 255 | |||
| 256 | ✗ | void add_edge(int from, int to, flow_t cap) { | |
| 257 | ✗ | return add_bi_edge(from, to, cap, 0); | |
| 258 | } | ||
| 259 | |||
| 260 | ✗ | void add_bi_edge(int from, int to, flow_t cap, flow_t rev_cap) { | |
| 261 | ✗ | assert(cap >= 0); | |
| 262 | ✗ | assert(rev_cap >= 0); | |
| 263 | ✗ | int e = int(edges.size()); | |
| 264 | ✗ | edges.emplace_back(edge_t{to, cap}); | |
| 265 | ✗ | edges.emplace_back(edge_t{from, rev_cap}); | |
| 266 | ✗ | adj[from].push_back(e); | |
| 267 | ✗ | adj[to].push_back(e+1); | |
| 268 | } | ||
| 269 | |||
| 270 | static constexpr tot_flow_t INF_FLOW = std::numeric_limits<tot_flow_t>::max() / 4; | ||
| 271 | std::vector<int> buf; | ||
| 272 | std::vector<int> level; | ||
| 273 | ✗ | tot_flow_t dinic_dfs(int cur, int t, tot_flow_t f) { | |
| 274 | ✗ | if (cur == t) return f; | |
| 275 | ✗ | tot_flow_t cur_f = 0; | |
| 276 | ✗ | assert(f > 0); | |
| 277 | ✗ | for (; buf[cur] < int(adj[cur].size()); buf[cur]++) { | |
| 278 | ✗ | int e = adj[cur][buf[cur]]; | |
| 279 | ✗ | int nxt = edges[e].dest; | |
| 280 | ✗ | if (level[nxt] == level[cur] + 1 && edges[e].cap > 0) { | |
| 281 | ✗ | flow_t v = flow_t(dinic_dfs(nxt, t, std::min<tot_flow_t>(f, edges[e].cap))); | |
| 282 | ✗ | edges[e].cap -= v; | |
| 283 | ✗ | edges[e^1].cap += v; | |
| 284 | ✗ | f -= v; | |
| 285 | ✗ | cur_f += v; | |
| 286 | ✗ | if (f == 0) break; | |
| 287 | } | ||
| 288 | } | ||
| 289 | ✗ | return cur_f; | |
| 290 | } | ||
| 291 | ✗ | tot_flow_t dinic(int s, int t) { | |
| 292 | ✗ | tot_flow_t tot_flow = 0; | |
| 293 | ✗ | while (true) { | |
| 294 | ✗ | buf.clear(); | |
| 295 | ✗ | buf.reserve(N); | |
| 296 | ✗ | level.assign(N, -1); | |
| 297 | ✗ | buf.push_back(s); | |
| 298 | ✗ | level[s] = 0; | |
| 299 | ✗ | for (int z = 0; z < int(buf.size()); z++) { | |
| 300 | ✗ | int cur = buf[z]; | |
| 301 | ✗ | for (int e : adj[cur]) { | |
| 302 | ✗ | int nxt = edges[e].dest; | |
| 303 | ✗ | if (edges[e].cap > 0 && level[nxt] == -1) { | |
| 304 | ✗ | level[nxt] = level[cur] + 1; | |
| 305 | ✗ | buf.push_back(nxt); | |
| 306 | } | ||
| 307 | } | ||
| 308 | } | ||
| 309 | ✗ | if (level[t] == -1) break; | |
| 310 | ✗ | buf.assign(N, 0); | |
| 311 | ✗ | tot_flow += dinic_dfs(s, t, INF_FLOW); | |
| 312 | } | ||
| 313 | ✗ | return tot_flow; | |
| 314 | } | ||
| 315 | ✗ | tot_flow_t max_flow(int s, int t) { return dinic(s, t); } | |
| 316 | }; | ||
| 317 |