leetcode

commit 852aa23f39a62fb346ce2e81c5f689b0f07af2c0
parent a8f78860b12f49e6df4c3532107f8c62828a6a48
Author: Dimitrije Dobrota <mail@dimitrijedobrota.com>
Date:   Sat, 24 Feb 2024 13:17:49 +0000

Improve UnionFind template

Diffstat:
M
Templates/Union_Find.cpp
 | 
30
++++++++++++++++++------------

1 file changed, 18 insertions(+), 12 deletions(-)
diff --git a/Templates/Union_Find.cpp b/Templates/Union_Find.cpp
@@ -1,28 +1,34 @@
 class UnionFind {
     int n, cnt = n;
-    vector<int> root, size;
+    mutable vector<int> root;
+    vector<int> size;
 
   public:
     UnionFind(int n) : n(n), root(n), size(n, 1) { iota(root.begin(), root.end(), 0); }
-
     UnionFind(const UnionFind &uf) : n(uf.n), cnt(uf.cnt), root(uf.root), size(uf.size) {}
 
-    int find(int x) {
-        while (x != root[x])
+    int find(int x) const {
+        while (x != root[x]) {
             x = root[x] = root[root[x]];
+        }
         return x;
     }
 
     void join(int x, int y) {
         x = find(x), y = find(y);
-        if (x != y) {
-            if (size[x] > size[y]) swap(x, y);
-            root[x] = y;
-            size[y] += size[x];
-            cnt--;
-        }
+        if (x == y) return;
+
+        if (size[x] > size[y]) swap(x, y);
+        root[x] = y;
+        size[y] += size[x];
+        cnt--;
+    }
+
+    void reset(int x) {
+        root[x] = x;
+        size[x] = 0;
     }
 
-    int count() { return cnt; }
-    bool connected(int x, int y) { return find(x) == find(y); }
+    int count() const { return cnt; }
+    bool connected(int x, int y) const { return find(x) == find(y); }
 };