bruih

Author: envyaims

solutions/gold/cf-468B.mdx

@@ -31,27 +31,27 @@ We can treat each set as a graph with several connected components. For each $p_
 #include <bits/stdc++.h>
 using namespace std;
 
+// BeginCodeSnip{Disjoint Set Union}
 struct DSU {
-	vector<int> p, sz;
-	DSU(int n) {
-		p.resize(n);
-		sz.resize(n, 1);
-		iota(p.begin(), p.end(), 0);
-	}
-	int get(int x) {
-		if (p[x] != x) { p[x] = get(p[x]); }
-		return p[x];
-	}
-	void unite(int a, int b) {
-		a = get(a);
-		b = get(b);
-		if (sz[a] < sz[b]) { swap(a, b); }
-		if (a != b) {
-			p[b] = a;
-			sz[a] += sz[b];
-		}
+	vector<int> e;
+	DSU(int N) { e = vector<int>(N, -1); }
+
+	// get representive component (uses path compression)
+	int get(int x) { return e[x] < 0 ? x : e[x] = get(e[x]); }
+	
+	bool same_set(int a, int b) { return get(a) == get(b); }
+	
+	int size(int x) { return -e[get(x)]; }
+	
+	bool unite(int x, int y) {  // union by size
+		x = get(x), y = get(y);
+		if (x == y) return false;
+		if (e[x] > e[y]) swap(x, y);
+		e[x] += e[y]; e[y] = x;
+		return true;
 	}
 };
+// EndCodeSnip
 
 int main() {
 	int n, a, b;