cp-library-cpp
This documentation is automatically generated by online-judge-tools/verification-helper
Assignment Problem
(library/graph/assignment_problem.hpp)
- View this file on GitHub
- Last update: 2023-01-02 17:33:27+09:00
- Include:
#include "library/graph/assignment_problem.hpp"
Assignment Problem
Verified with
Code
#ifndef SUISEN_MIN_COST_FLOW
#define SUISEN_MIN_COST_FLOW
#include <algorithm>
#include <cassert>
#include <queue>
#include <limits>
#include <vector>
namespace suisen {
template <typename T>
std::pair<T, std::vector<int>> assignment_problem(const std::vector<std::vector<T>>& a) {
static constexpr T INF_COST = std::numeric_limits<T>::max() / 2;
const int n = a.size(), k = 2 * n + 2;
const int s = 2 * n, t = s + 1;
struct Edge {
int to;
int cap;
T cost;
int rev;
};
std::vector<std::vector<Edge>> g(k);
std::vector<T> potential(k);
std::vector<T> dist(k);
std::vector<int> prev_vid(k), prev_eid(k);
std::vector<std::pair<int, int>> edges;
auto add_edge = [&](int u, int v, int cap, T cost) {
edges.emplace_back(u, g[u].size());
g[u].push_back(Edge { v, cap, cost, int(g[v].size()) });
g[v].push_back(Edge { u, 0, -cost, int(g[u].size()) - 1 });
};
for (int i = 0; i < n; ++i) {
add_edge(s, i, 1, 0);
for (int j = 0; j < n; ++j) {
add_edge(i, n + j, 1, a[i][j]);
}
add_edge(n + i, t, 1, 0);
}
auto update_dist = [&](int u, int eid) {
if (dist[u] == INF_COST) return false;
const auto& e = g[u][eid];
if (e.cap == 0) return false;
const int v = e.to;
T cost = e.cost + potential[u] - potential[v];
if (dist[u] + cost < dist[v]) {
dist[v] = dist[u] + cost;
prev_vid[v] = u;
prev_eid[v] = eid;
return true;
}
return false;
};
auto sssp_dijkstra = [&]() {
std::vector<int8_t> used(k, false);
dist.assign(k, INF_COST);
dist[s] = 0;
while (true) {
T min_dist = INF_COST;
int arg_min = -1;
for (int i = 0; i < k; ++i) if (not used[i]) {
if (dist[i] < min_dist) {
min_dist = dist[i];
arg_min = i;
}
}
const int u = arg_min;
if (u == -1) return;
for (int i = 0; i < int(g[u].size()); ++i) {
const auto& e = g[u][i];
if (e.cap == 0) continue;
update_dist(u, i);
}
used[u] = true;
}
};
auto sssp_dag = [&]() {
std::vector<int> in(k, 0);
for (int u = 0; u < k; ++u) {
for (const auto& e : g[u]) {
if (e.cap == 0) continue;
++in[e.to];
}
}
std::deque<int> dq;
for (int u = 0; u < k; ++u) {
if (in[u] == 0) dq.push_back(u);
}
dist.assign(k, INF_COST);
dist[s] = 0;
while (dq.size()) {
int u = dq.front();
dq.pop_front();
for (int i = 0; i < int(g[u].size()); ++i) {
const auto& e = g[u][i];
if (e.cap == 0) continue;
update_dist(u, i);
if (--in[e.to] == 0) {
dq.push_back(e.to);
}
}
}
};
auto update_potential = [&]() {
for (int u = 0; u < k; ++u) {
if (potential[u] != INF_COST) potential[u] += dist[u];
}
};
sssp_dag();
update_potential();
T cost = 0;
while (dist[t] != INF_COST) {
int df = n;
for (int v = t; v != s; v = prev_vid[v]) {
df = std::min(df, g[prev_vid[v]][prev_eid[v]].cap);
}
cost += df * potential[t];
for (int v = t; v != s; v = prev_vid[v]) {
auto& e = g[prev_vid[v]][prev_eid[v]];
e.cap -= df;
g[v][e.rev].cap += df;
}
sssp_dijkstra();
update_potential();
}
std::vector<int> assignment(n);
for (int i = 0; i < n; ++i) {
for (const Edge &e : g[i]) {
int j = e.to - n;
if (j >= n or e.cap == 1) continue;
assignment[i] = j;
}
}
return { cost, assignment };
};
} // namespace suisen
#endif // SUISEN_MIN_COST_FLOW#line 1 "library/graph/assignment_problem.hpp"
#include <algorithm>
#include <cassert>
#include <queue>
#include <limits>
#include <vector>
namespace suisen {
template <typename T>
std::pair<T, std::vector<int>> assignment_problem(const std::vector<std::vector<T>>& a) {
static constexpr T INF_COST = std::numeric_limits<T>::max() / 2;
const int n = a.size(), k = 2 * n + 2;
const int s = 2 * n, t = s + 1;
struct Edge {
int to;
int cap;
T cost;
int rev;
};
std::vector<std::vector<Edge>> g(k);
std::vector<T> potential(k);
std::vector<T> dist(k);
std::vector<int> prev_vid(k), prev_eid(k);
std::vector<std::pair<int, int>> edges;
auto add_edge = [&](int u, int v, int cap, T cost) {
edges.emplace_back(u, g[u].size());
g[u].push_back(Edge { v, cap, cost, int(g[v].size()) });
g[v].push_back(Edge { u, 0, -cost, int(g[u].size()) - 1 });
};
for (int i = 0; i < n; ++i) {
add_edge(s, i, 1, 0);
for (int j = 0; j < n; ++j) {
add_edge(i, n + j, 1, a[i][j]);
}
add_edge(n + i, t, 1, 0);
}
auto update_dist = [&](int u, int eid) {
if (dist[u] == INF_COST) return false;
const auto& e = g[u][eid];
if (e.cap == 0) return false;
const int v = e.to;
T cost = e.cost + potential[u] - potential[v];
if (dist[u] + cost < dist[v]) {
dist[v] = dist[u] + cost;
prev_vid[v] = u;
prev_eid[v] = eid;
return true;
}
return false;
};
auto sssp_dijkstra = [&]() {
std::vector<int8_t> used(k, false);
dist.assign(k, INF_COST);
dist[s] = 0;
while (true) {
T min_dist = INF_COST;
int arg_min = -1;
for (int i = 0; i < k; ++i) if (not used[i]) {
if (dist[i] < min_dist) {
min_dist = dist[i];
arg_min = i;
}
}
const int u = arg_min;
if (u == -1) return;
for (int i = 0; i < int(g[u].size()); ++i) {
const auto& e = g[u][i];
if (e.cap == 0) continue;
update_dist(u, i);
}
used[u] = true;
}
};
auto sssp_dag = [&]() {
std::vector<int> in(k, 0);
for (int u = 0; u < k; ++u) {
for (const auto& e : g[u]) {
if (e.cap == 0) continue;
++in[e.to];
}
}
std::deque<int> dq;
for (int u = 0; u < k; ++u) {
if (in[u] == 0) dq.push_back(u);
}
dist.assign(k, INF_COST);
dist[s] = 0;
while (dq.size()) {
int u = dq.front();
dq.pop_front();
for (int i = 0; i < int(g[u].size()); ++i) {
const auto& e = g[u][i];
if (e.cap == 0) continue;
update_dist(u, i);
if (--in[e.to] == 0) {
dq.push_back(e.to);
}
}
}
};
auto update_potential = [&]() {
for (int u = 0; u < k; ++u) {
if (potential[u] != INF_COST) potential[u] += dist[u];
}
};
sssp_dag();
update_potential();
T cost = 0;
while (dist[t] != INF_COST) {
int df = n;
for (int v = t; v != s; v = prev_vid[v]) {
df = std::min(df, g[prev_vid[v]][prev_eid[v]].cap);
}
cost += df * potential[t];
for (int v = t; v != s; v = prev_vid[v]) {
auto& e = g[prev_vid[v]][prev_eid[v]];
e.cap -= df;
g[v][e.rev].cap += df;
}
sssp_dijkstra();
update_potential();
}
std::vector<int> assignment(n);
for (int i = 0; i < n; ++i) {
for (const Edge &e : g[i]) {
int j = e.to - n;
if (j >= n or e.cap == 1) continue;
assignment[i] = j;
}
}
return { cost, assignment };
};
} // namespace suisen