MST + test

Author: adamant-pwn

cp-algo/data_structures/dsu.hpp

@@ -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

cp-algo/graph/base.hpp

@@ -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();}

cp-algo/graph/euler.hpp

@@ -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);
@@ -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);
                     }
                 }

cp-algo/graph/mst.hpp

@@ -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

verify/graph/mst.test.cpp

@@ -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();
+    }
+}