add dsu

Author: lrvideckis

library/data_structures/dsu/dsu.hpp

@@ -0,0 +1,12 @@
+#pragma once
+struct dsu {
+  vi e;
+  dsu(int n): e(n, -1) {}
+  int size(int x) { return -e[go(x)]; }
+  int go(int x) { return e[x] < 0 ? x : e[x] = go(e[x]); }
+  bool join(int a, int b) {
+    if ((a = go(a)) == (b = go(b))) return 0;
+    if (e[a] > e[b]) swap(a, b);
+    return e[a] += e[b], e[b] = a, 1;
+  }
+};

library/graphs/mst.hpp

@@ -1,5 +1,5 @@
 #pragma once
-#include "../../kactl/content/data-structures/UnionFind.h"
+#include "../data_structures/dsu/dsu.hpp"
 //! @code
 //!   auto [mst_cost, ids] = mst(w_eds, n);
 //! @endcode
@@ -14,7 +14,7 @@ pair<ll, vi> mst(const vector<array<int, 3>>& w_eds,
   ranges::sort(order, [&](int i, int j) {
     return w_eds[i][2] < w_eds[j][2];
   });
-  UF uf(n);
+  dsu uf(n);
   vi ids;
   ll cost = 0;
   for (int it : order) {

tests/library_checker_aizu_tests/data_structures/dsu.test.cpp

@@ -0,0 +1,16 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/unionfind"
+#include "../template.hpp"
+#include "../../../library/data_structures/dsu/dsu.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n, q;
+  cin >> n >> q;
+  dsu uf(n);
+  while (q--) {
+    int type, u, v;
+    cin >> type >> u >> v;
+    if (type == 0) uf.join(u, v);
+    else cout << (uf.go(u) == uf.go(v)) << '\n';
+  }
+  return 0;
+}

tests/library_checker_aizu_tests/graphs/bcc_callback_lib_checker_two_cc.test.cpp

@@ -2,7 +2,7 @@
   "https://judge.yosupo.jp/problem/two_edge_connected_components"
 #include "../template.hpp"
 #include "../../../library/graphs/bcc_callback.hpp"
-#include "../../../kactl/content/data-structures/UnionFind.h"
+#include "../../../library/data_structures/dsu/dsu.hpp"
 int main() {
   cin.tie(0)->sync_with_stdio(0);
   int n, m;
@@ -14,7 +14,7 @@ int main() {
     adj[u].push_back(v);
     adj[v].push_back(u);
   }
-  UF uf(n);
+  dsu uf(n);
   vector<bool> seen(n);
   bcc_callback(adj, [&](const vi& nodes) {
     int count_edges = 0;
@@ -31,7 +31,7 @@ int main() {
     for (int v : nodes) uf.join(v, nodes[0]);
   });
   vector<vi> two_edge_ccs(n);
-  rep(i, 0, n) two_edge_ccs[uf.find(i)].push_back(i);
+  rep(i, 0, n) two_edge_ccs[uf.go(i)].push_back(i);
   int cnt_ccs = 0;
   rep(i, 0, n) cnt_ccs += (!empty(two_edge_ccs[i]));
   cout << cnt_ccs << '\n';

tests/library_checker_aizu_tests/graphs/offline_incremental_scc.test.cpp

@@ -2,7 +2,7 @@
   "https://judge.yosupo.jp/problem/incremental_scc"
 #include "../template.hpp"
 #include "../../../library/graphs/strongly_connected_components/offline_incremental_scc.hpp"
-#include "../../../kactl/content/data-structures/UnionFind.h"
+#include "../../../library/data_structures/dsu/dsu.hpp"
 #include "../../../library/math/mod_int.hpp"
 int main() {
   cin.tie(0)->sync_with_stdio(0);
@@ -20,13 +20,13 @@ int main() {
   iota(all(order), 0);
   ranges::sort(all(order), {},
     [&](int i) { return joins[i]; });
-  UF uf(n);
+  dsu uf(n);
   mint sum = 0;
   for (int t = 0, it = 0; t < m; t++) {
     while (it < m && joins[order[it]] <= t) {
       auto [u, v] = eds[order[it]];
-      u = uf.find(u);
-      v = uf.find(v);
+      u = uf.go(u);
+      v = uf.go(v);
       if (uf.e[u] > uf.e[v]) swap(u, v);
       if (u != v) {
         sum = sum + xs[u] * xs[v];

tests/library_checker_aizu_tests/handmade_tests/dsu.test.cpp

@@ -0,0 +1,51 @@
+#define PROBLEM \
+  "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
+#include "../template.hpp"
+#include "../../../library/contest/random.hpp"
+#include "../../../library/data_structures/dsu/dsu.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  for (int n = 0; n <= 100; n++) {
+    vector<vector<int>> adj(n);
+    dsu uf(n);
+    for (int q = 0; q < n; q++) {
+      int u = rnd(0, n - 1);
+      int v = rnd(0, n - 1);
+      bool joined = uf.join(u, v);
+      {
+        vector<bool> vis(n);
+        auto dfs = [&](auto&& self, int v) -> void {
+          for (int u : adj[v]) {
+            if (!vis[u]) {
+              vis[u] = 1;
+              self(self, u);
+            }
+          }
+        };
+        vis[u] = 1;
+        dfs(dfs, u);
+        assert(vis[v] != joined);
+      }
+      adj[u].push_back(v);
+      adj[v].push_back(u);
+      for (int i = 0; i < n; i++) {
+        vector<bool> vis(n);
+        int cnt_nodes = 0;
+        auto dfs = [&](auto&& self, int v) -> void {
+          cnt_nodes++;
+          for (int u : adj[v]) {
+            if (!vis[u]) {
+              vis[u] = 1;
+              self(self, u);
+            }
+          }
+        };
+        vis[i] = 1;
+        dfs(dfs, i);
+        assert(cnt_nodes == uf.size(i));
+      }
+    }
+  }
+  cout << "Hello World\n";
+  return 0;
+}