Add UnionFindStore (#3)

content/data-structures/UnionFindStore.h

@@ -0,0 +1,29 @@
+/**
+ * Author: Pietro Palombini
+ * Date: 2024-03-01
+ * License: CC0
+ * Source: folklore
+ * Description: Disjoint-set data structure, with support for
+ * storing additional data in each set.
+ * Time: $O(\alpha(N))$
+ * Status: tested on CSES: Road Construction
+ */
+#pragma once
+
+typedef ll D;
+struct UFStore {
+	void merge(D& large, const D& small) { large += small; }
+	vi e;
+	vector<D> d;
+	UFStore(ll n) : e(n, -1), d(n) {}
+	UFStore(vector<D>& d) : e(SZ(d), -1), d(d) {}
+	D& find(ll x) { return d[repr(x)]; }
+	ll repr(ll x) { return e[x] < 0 ? x : e[x] = repr(e[x]); }
+	bool join(ll a, ll b) {
+		a = repr(a), b = repr(b);
+		if (a == b) return false;
+		if (e[a] > e[b]) swap(a, b);
+		e[a] += e[b], e[b] = a, merge(d[a], d[b]);
+		return true;
+	}
+};

content/data-structures/chapter.tex

@@ -6,6 +6,7 @@ \chapter{Data structures}
 \kactlimport{LazySegmentTree.h}
 \kactlimport{MemoryLazySegmentTree.h}
 \kactlimport{UnionFind.h}
+\kactlimport{UnionFindStore.h}
 \kactlimport{UnionFindRollback.h}
 \kactlimport{SubMatrix.h}
 \kactlimport{Matrix.h}

test-problems/data-structures/UnionFindStore.cpp

@@ -0,0 +1,50 @@
+// Problem: https://cses.fi/problemset/task/1676
+// Submission: https://cses.fi/problemset/result/8650362/
+#include <bits/stdc++.h>
+using namespace std;
+#define fst first
+#define snd second
+#define pb push_back
+#define fore(i, a, gmat) for (ll i = a; i < gmat; i++)
+#define ALL(x) begin(x), end(x)
+#define SZ(x) (ll)(x).size()
+typedef long long ll;
+typedef pair<ll, ll> ii;
+typedef vector<ll> vi;
+
+struct D {
+	ll sz = 1;
+};
+struct UFStore {
+	void merge(D& large, const D& small) { large.sz += small.sz; }
+	vi e;
+	vector<D> d;
+	UFStore(ll n) : e(n, -1), d(n) {}
+	UFStore(vector<D>& d) : e(SZ(d), -1), d(d) {}
+	D& find(ll x) { return d[repr(x)]; }
+	ll repr(ll x) { return e[x] < 0 ? x : e[x] = repr(e[x]); }
+	bool join(ll a, ll b) {
+		a = repr(a), b = repr(b);
+		if (a == b) return false;
+		if (e[a] > e[b]) swap(a, b);
+		e[a] += e[b], e[b] = a, merge(d[a], d[b]);
+		return true;
+	}
+};
+
+int main() {
+	cin.tie(0)->sync_with_stdio(0);
+
+	ll n, m;
+	cin >> n >> m;
+	ll largest = 1, cnt = n;
+	UFStore uf(n);
+	fore(i, 0, m) {
+		ll a, b;
+		cin >> a >> b;
+		a--, b--;
+		if (uf.join(a, b)) cnt--;
+		largest = max(largest, uf.find(a).sz);
+		cout << cnt << " " << largest << "\n";
+	}
+}