Skip to content

Commit

Permalink
MST + test
Browse files Browse the repository at this point in the history
  • Loading branch information
adamant-pwn committed Aug 6, 2024
1 parent f91796a commit fec4131
Show file tree
Hide file tree
Showing 5 changed files with 98 additions and 11 deletions.
24 changes: 24 additions & 0 deletions cp-algo/data_structures/dsu.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
#ifndef CP_ALGO_DATA_STRUCTURES_DSU_HPP
#define CP_ALGO_DATA_STRUCTURES_DSU_HPP
#include <numeric>
#include <vector>
namespace cp_algo::data_structures {
struct disjoint_set_union {
disjoint_set_union (int n): par(n) {
std::iota(begin(par), end(par), 0);
}
int get(int v) {
return v == par[v] ? v : par[v] = get(par[v]);
}
bool uni(int a, int b) {
a = get(a);
b = get(b);
par[a] = b;
return a != b;
}
private:
std::vector<int> par;
};
using dsu = disjoint_set_union;
}
#endif // CP_ALGO_DATA_STRUCTURES_DSU_HPP
10 changes: 8 additions & 2 deletions cp-algo/graph/base.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,19 @@ namespace cp_algo::graph {
call_adjacent(v, callback, [](){return false;});
}
void call_edges(auto &&callback) const {
for(node_index v: nodes_view()) {
call_adjacent(v, [&](edge_index e) {callback(v, e);});
for(edge_index e: edges_view()) {
callback(e);
}
}
auto nodes_view() const {
return std::views::iota(0, n());
}
auto edges_view() const {
return std::views::filter(
std::views::iota(0, 2 * m()),
[](edge_index e) {return !(e % 2);}
);
}
auto const& incidence_lists() const {return _adj;}
edge_t const& edge(edge_index e) const {return edges[e];}
node_index n() const {return _adj.size();}
Expand Down
20 changes: 11 additions & 9 deletions cp-algo/graph/euler.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,18 @@ namespace cp_algo::graph {
std::vector<int> deg(g.n());
constexpr bool undirected = graph::undirected;
int res = 0;
g.call_edges([&](int u, int e) {
g.call_edges([&](edge_index e) {
int u = g.edge(e ^ 1).to;
int v = g.edge(e).to;
res = u;
if constexpr (undirected) {
deg[u] ^= 1;
} else {
deg[u]++;
deg[g.edge(e).to]--;
}
deg[u]++;
deg[v]--;
});
if constexpr (undirected) {
for(auto &it: deg) {
it = bool(it % 2);
}
}
auto nodes = g.nodes_view();
auto is_start = [&](int v) {return deg[v] > 0;};
auto starts = std::ranges::count_if(nodes, is_start);
Expand All @@ -42,8 +45,7 @@ namespace cp_algo::graph {
int e = adj.data[std::exchange(head[v], adj.next[head[v]])];
if(!used[e / 2]) {
used[e / 2] = 1;
int u = g.edge(e).to;
self(self, u);
self(self, g.edge(e).to);
trail.push_back(e);
}
}
Expand Down
26 changes: 26 additions & 0 deletions cp-algo/graph/mst.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
#ifndef CP_ALGO_GRAPH_MST_HPP
#define CP_ALGO_GRAPH_MST_HPP
#include "base.hpp"
#include "../data_structures/dsu.hpp"
#include <algorithm>
namespace cp_algo::graph {
template<weighted_edge_type edge_t>
auto mst(graph<undirected, edge_t> const& g) {
std::vector<std::pair<int64_t, edge_index>> edges;
g.call_edges([&](edge_index e) {
edges.emplace_back(g.edge(e).w, e);
});
std::ranges::sort(edges);
data_structures::dsu me(g.n());
int64_t total = 0;
std::vector<edge_index> mst;
for(auto [w, e]: edges) {
if(me.uni(g.edge(e ^ 1).to, g.edge(e).to)) {
total += w;
mst.push_back(e);
}
}
return std::pair{total, mst};
}
}
#endif // CP_ALGO_GRAPH_MST_HPP
29 changes: 29 additions & 0 deletions verify/graph/mst.test.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
// @brief Minimum Spanning Tree
#define PROBLEM "https://judge.yosupo.jp/problem/minimum_spanning_tree"
#pragma GCC optimize("Ofast,unroll-loops")
#include "cp-algo/graph/mst.hpp"
#include <bits/stdc++.h>

using namespace std;
using namespace cp_algo::graph;

void solve() {
int n, m;
cin >> n >> m;
graph<undirected, weighted_edge> g(n);
g.read_edges(m);
auto [X, E] = mst(g);
cout << X << "\n";
for(int e: E) {cout << e / 2 << " ";}
}

signed main() {
//freopen("input.txt", "r", stdin);
ios::sync_with_stdio(0);
cin.tie(0);
int t = 1;
//cin >> t;
while(t--) {
solve();
}
}

0 comments on commit fec4131

Please sign in to comment.