MST: make code more reusable + add test

Graph/mst.h

@@ -9,32 +9,32 @@
 
 // MST {{{
 using ll = long long;
-struct Edge {
-    int u, v;
-    ll c;
-};
-bool operator < (const Edge& a, const Edge& b) {
-    return a.c < b.c;
-}
-ostream& operator << (ostream& out, const Edge& e) {
-    out << e.u << " - " << e.v << " [" << e.c << ']';
-    return out;
-}
-std::pair<ll, std::vector<Edge>> mst(
+template<typename EdgeT>
+std::pair<ll, std::vector<EdgeT>> mst(
         int n,
-        std::vector<Edge> edges) {
+        std::vector<EdgeT> edges) {
     std::sort(edges.begin(), edges.end());
 
     DSU dsu(n + 1);  // tolerate 1-based index
     ll total = 0;
-    vector<Edge> tree;
+    vector<EdgeT> tree;
     for (const auto& e : edges) {
-        const auto [u, v, c] = e;
-        if (dsu.merge(u, v)) {
-            total += c;
+        if (dsu.merge(e.u, e.v)) {
+            total += e.c;
             tree.push_back(e);
         }
     }
     return {total, tree};
 }
+struct Edge {
+    int u, v;
+    ll c;
+};
+bool operator < (const Edge& a, const Edge& b) {
+    return a.c < b.c;
+}
+ostream& operator << (ostream& out, const Edge& e) {
+    out << e.u << " - " << e.v << " [" << e.c << ']';
+    return out;
+}
 // }}}

Graph/tests/yosupo_mst.test.cpp

@@ -0,0 +1,22 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/minimum_spanning_tree"
+
+#include "../../template.h"
+#include "../mst.h"
+
+struct E : Edge {
+    int id;
+};
+
+void solve() {
+    int n, m; cin >> n >> m;
+    vector<E> edges(m);
+    REP(i,m) {
+        auto& e = edges[i];
+        cin >> e.u >> e.v >> e.c;
+        e.id = i;
+    }
+    auto g = mst<E>(n, edges);
+    cout << g.first << '\n';
+    for (auto& e : g.second) cout << e.id << ' ';
+    cout << '\n';
+}