.

include/utils.hpp

@@ -1,12 +1,13 @@
 #pragma once
 
-#include <endian.h>
 #include <vector>
 #include <unordered_map>
 
 namespace utils {
 
-using IdType = std::size_t;
+using std::size_t;
+
+using IdType = size_t;
 
 enum class ReferenceType { Reference, UniqueReference };
 
@@ -39,4 +40,42 @@ private:
   std::unordered_map<T, IdType> value_to_id_;
 };
 
+class Union { // TODO: recall right algorithm name
+public:
+  Union(size_t n) {
+    edges_.resize(n);
+    ranks_.resize(n);
+
+    for (size_t i = 0; i < n; ++i) {
+      edges_[i] = i;
+      ranks_[i] = 1;
+    }
+  }
+
+  void Connect(size_t u, size_t v) { // TODO: recall choice algorithm
+    u = GetRoot(u);
+    v = GetRoot(v);
+    if (ranks_[v] >= ranks_[u]) {
+      edges_[u] = v;
+      ranks_[v] = std::max(ranks_[u] + 1, ranks_[v]);
+    } else {
+      edges_[v] = u;
+    }
+  }
+
+  bool IsConnected(size_t u, size_t v) {
+    return GetRoot(u) == GetRoot(v);
+  }
+
+  size_t GetRoot(size_t v) {
+    if (edges_[v] == v) {
+      return v;
+    }
+    return edges_[v] = GetRoot(edges_[v]);
+  }
+private:
+  std::vector<size_t> edges_;
+  std::vector<size_t> ranks_;
+};
+
 } // namespace utils