test/AOJ/GRL_2_B.test.cpp
Depends on
Code
#define PROBLEM \
"https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_2_B"
#include <bits/stdc++.h>
#include "library/graph/MinimumSpanningArborescence.hpp"
#include "library/graph/WeightedGraph.hpp"
int main() {
int n, m, r;
std::cin >> n >> m >> r;
WeightedGraph<int> g(n, m, true, 0);
auto ans = minimum_spanning_arborescence(g, r);
assert(ans.has_value());
std::cout << ans.value().first << std::endl;
}
#line 1 "test/AOJ/GRL_2_B.test.cpp"
#define PROBLEM \
"https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_2_B"
#include <bits/stdc++.h>
#line 3 "library/datastructure/unionfind/UnionFind.hpp"
class UnionFind {
int n, num;
std::vector<int> sz, parent;
public:
UnionFind() = default;
UnionFind(int n) : n(n), num(n), sz(n, 1), parent(n, 0) {
std::iota(parent.begin(), parent.end(), 0);
}
int leader(int x) {
assert(0 <= x and x < n);
return (x == parent[x] ? x : parent[x] = leader(parent[x]));
}
bool same(int x, int y) {
assert(0 <= x and x < n and 0 <= y and y < n);
return leader(x) == leader(y);
}
bool merge(int x, int y) {
assert(0 <= x and x < n and 0 <= y and y < n);
x = leader(x);
y = leader(y);
if (x == y)
return false;
if (sz[x] < sz[y])
std::swap(x, y);
sz[x] += sz[y];
parent[y] = x;
num--;
return true;
}
int size(const int x) {
assert(0 <= x and x < n);
return sz[leader(x)];
}
int count() const { return num; }
std::vector<std::vector<int>> groups() {
std::vector<std::vector<int>> res(n);
for (int i = 0; i < n; i++)
res[leader(i)].push_back(i);
std::erase_if(res, [](const auto &vec) { return vec.empty(); });
return res;
}
};
#line 3 "library/graph/MinimumSpanningArborescence.hpp"
template <typename WG, typename W = typename WG::weight_type>
std::optional<std::pair<W, std::vector<int>>>
minimum_spanning_arborescence(WG g, int r = 0) {
int n = g.n;
W res = 0;
std::vector<W> new_add(n, 0);
std::vector<int> tree(n), pre(n), state(n, 0);
UnionFind uf(n);
state[r] = 2;
auto compare = [&](const int &a, const int &b) {
return g.edges[a].weight > g.edges[b].weight;
};
using PQ = std::priority_queue<int, std::vector<int>, decltype(compare)>;
std::vector<std::pair<PQ, W>> pq_add(n, {PQ{compare}, 0});
for (int i = 0; i < g.edges.size(); i++)
pq_add[g.edges[i].to].first.push(i);
std::vector<int> pq_id(n);
std::iota(pq_id.begin(), pq_id.end(), 0);
auto merge = [&](int u, int v) {
u = uf.leader(u);
v = uf.leader(v);
if (u == v)
return;
uf.merge(u, v);
auto &[pq1, add1] = pq_add[pq_id[u]];
auto &[pq2, add2] = pq_add[pq_id[v]];
if (pq1.size() > pq2.size()) {
while (pq2.size()) {
int edge_id = pq2.top();
pq2.pop();
g.edges[edge_id].weight -= add2 - add1;
pq1.push(edge_id);
}
pq_id[uf.leader(v)] = pq_id[u];
} else {
while (pq1.size()) {
int edge_id = pq1.top();
pq1.pop();
g.edges[edge_id].weight -= add1 - add2;
pq2.push(edge_id);
}
pq_id[uf.leader(v)] = pq_id[v];
}
};
for (int i = 0; i < n; i++) {
int now = uf.leader(i);
if (state[now])
continue;
std::vector<int> processing;
while (state[now] != 2) {
processing.push_back(now);
state[now] = 1;
auto &[pq, add] = pq_add[pq_id[now]];
if (!pq.size())
return std::nullopt;
int edge_id = pq.top();
pq.pop();
auto &e = g.edges[edge_id];
res += e.weight - add;
tree[e.to] = edge_id;
pre[now] = uf.leader(e.from);
new_add[now] = e.weight;
if (state[pre[now]] == 1) {
int v = now;
do {
pq_add[pq_id[v]].second = new_add[v];
merge(v, now);
v = uf.leader(pre[v]);
} while (!uf.same(v, now));
now = uf.leader(now);
} else
now = uf.leader(pre[now]);
}
for (int v : processing)
state[v] = 2;
}
tree.erase(tree.begin() + r);
return std::make_pair(res, tree);
}
#line 2 "library/graph/WeightedGraph.hpp"
template <typename T> struct WeightedEdge {
WeightedEdge() = default;
WeightedEdge(int from, int to, T weight)
: from(from), to(to), weight(weight) {}
int from, to;
T weight;
operator int() const { return to; }
};
template <typename T> struct WeightedGraph {
int n;
using weight_type = T;
using edge_type = WeightedEdge<T>;
std::vector<edge_type> edges;
protected:
std::vector<int> in_deg;
bool prepared;
class OutgoingEdges {
WeightedGraph *g;
int l, r;
public:
OutgoingEdges(WeightedGraph *g, int l, int r) : g(g), l(l), r(r) {}
edge_type *begin() { return &(g->edges[l]); }
edge_type *end() { return &(g->edges[r]); }
edge_type &operator[](int i) { return g->edges[l + i]; }
int size() const { return r - l; }
};
class ConstOutgoingEdges {
const WeightedGraph *g;
int l, r;
public:
ConstOutgoingEdges(const WeightedGraph *g, int l, int r)
: g(g), l(l), r(r) {}
const edge_type *begin() const { return &(g->edges[l]); }
const edge_type *end() const { return &(g->edges[r]); }
const edge_type &operator[](int i) const { return g->edges[l + i]; }
int size() const { return r - l; }
};
public:
OutgoingEdges operator[](int v) {
assert(prepared);
return {this, in_deg[v], in_deg[v + 1]};
}
const ConstOutgoingEdges operator[](int v) const {
assert(prepared);
return {this, in_deg[v], in_deg[v + 1]};
}
bool is_prepared() const { return prepared; }
WeightedGraph() : n(0), in_deg(1, 0), prepared(false) {}
WeightedGraph(int n) : n(n), in_deg(n + 1, 0), prepared(false) {}
WeightedGraph(int n, int m, bool directed = false, int indexed = 1)
: n(n), in_deg(n + 1, 0), prepared(false) {
scan(m, directed, indexed);
}
void resize(int n) { n = n; }
void add_arc(int from, int to, T weight) {
assert(!prepared);
assert(0 <= from and from < n and 0 <= to and to < n);
edges.emplace_back(from, to, weight);
in_deg[from + 1]++;
}
void add_edge(int u, int v, T weight) {
add_arc(u, v, weight);
add_arc(v, u, weight);
}
void add_arc(const edge_type &e) { add_arc(e.from, e.to, e.weight); }
void add_edge(const edge_type &e) { add_edge(e.from, e.to, e.weight); }
void scan(int m, bool directed = false, int indexed = 1) {
edges.reserve(directed ? m : 2 * m);
while (m--) {
int u, v;
std::cin >> u >> v;
u -= indexed;
v -= indexed;
T weight;
std::cin >> weight;
if (directed)
add_arc(u, v, weight);
else
add_edge(u, v, weight);
}
build();
}
void build() {
assert(!prepared);
prepared = true;
for (int v = 0; v < n; v++)
in_deg[v + 1] += in_deg[v];
std::vector<edge_type> new_edges(in_deg.back());
auto counter = in_deg;
for (auto &&e : edges)
new_edges[counter[e.from]++] = e;
edges = new_edges;
}
void graph_debug() const {
#ifndef __DEBUG
return;
#endif
assert(prepared);
for (int from = 0; from < n; from++) {
std::cerr << from << ";";
for (int i = in_deg[from]; i < in_deg[from + 1]; i++)
std::cerr << "(" << edges[i].to << "," << edges[i].weight
<< ")";
std::cerr << "\n";
}
}
};
#line 7 "test/AOJ/GRL_2_B.test.cpp"
int main() {
int n, m, r;
std::cin >> n >> m >> r;
WeightedGraph<int> g(n, m, true, 0);
auto ans = minimum_spanning_arborescence(g, r);
assert(ans.has_value());
std::cout << ans.value().first << std::endl;
}
Back to top page
build/pch/stdc++.hpp