feat: Implement Union-Find data structure and corresponding test cases

Author: weilycoder

test/unionfind.test.cpp

@@ -0,0 +1,27 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/unionfind"
+
+#include "weilycoder/unionfind.hpp"
+#include <iostream>
+using namespace std;
+using namespace weilycoder;
+
+int main() {
+  cin.tie(nullptr)->sync_with_stdio(false);
+  cin.exceptions(cin.failbit | cin.badbit);
+  size_t n, q;
+  cin >> n >> q;
+  UnionFind uf(n);
+  for (size_t _ = 0; _ < q; ++_) {
+    size_t op, x, y;
+    cin >> op >> x >> y;
+    switch (op) {
+    case 0:
+      uf.unite(x, y);
+      break;
+    case 1:
+      cout << (uf.getf(x) == uf.getf(y) ? "1\n" : "0\n");
+      break;
+    }
+  }
+  return 0;
+}

weilycoder/unionfind.hpp

@@ -0,0 +1,29 @@
+#include <algorithm>
+#include <numeric>
+#include <vector>
+
+namespace weilycoder {
+template <typename ptr_t = size_t> struct UnionFind {
+  std::vector<ptr_t> parent, size;
+
+  UnionFind(ptr_t n) : parent(n), size(n, 1) {
+    std::iota(parent.begin(), parent.end(), 0);
+  }
+
+  ptr_t getf(ptr_t x) {
+    while (x != parent[x])
+      x = parent[x] = parent[parent[x]];
+    return x;
+  }
+
+  bool unite(ptr_t x, ptr_t y) {
+    x = getf(x), y = getf(y);
+    if (x == y)
+      return false;
+    if (size[x] < size[y])
+      std::swap(x, y);
+    parent[y] = x, size[x] += size[y];
+    return true;
+  }
+};
+} // namespace weilycoder