From f31a486e6c42a2e3099d2b73d7260373f9326fa3 Mon Sep 17 00:00:00 2001
From: Trung Nguyen <ngthanhtrung23@gmail.com>
Date: Sun, 17 Mar 2024 02:51:32 +0000
Subject: [PATCH] MST: make code more reusable + add test

---
 Graph/mst.h                     | 34 ++++++++++++++++-----------------
 Graph/tests/yosupo_mst.test.cpp | 22 +++++++++++++++++++++
 2 files changed, 39 insertions(+), 17 deletions(-)
 create mode 100644 Graph/tests/yosupo_mst.test.cpp

diff --git a/Graph/mst.h b/Graph/mst.h
index 012db7d..3ff3eec 100644
--- a/Graph/mst.h
+++ b/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;
+}
 // }}}
diff --git a/Graph/tests/yosupo_mst.test.cpp b/Graph/tests/yosupo_mst.test.cpp
new file mode 100644
index 0000000..31d167d
--- /dev/null
+++ b/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';
+}