make union find a struct, remove kruskal

4 files changed, 50 insertions, 39 deletions

diff --git a/test/datastructures/unionFind.cpp b/test/datastructures/unionFind.cpp
index 2afdc86..4783f6b 100644
--- a/test/datastructures/unionFind.cpp
+++ b/test/datastructures/unionFind.cpp
@@ -1,8 +1,5 @@
 #include "../util.h"
-struct UF {
-	UF(int n) {init(n);}
-	#include <datastructures/unionFind.cpp>
-};
+#include <datastructures/unionFind.cpp>
 
 struct Naive {
 	vector<vector<int>> adj;
@@ -28,15 +25,18 @@ struct Naive {
 		}
 	}
 
-	int findSet(int a) {
+	int find(int a) {
 		int res = a;
 		dfs(a, [&](int x){res = min(res, x);});
 		return res;
 	}
 
-	void unionSets(int a, int b) {
+	bool link(int a, int b) {
+		bool linked = false;
+		dfs(a, [&](int x) { linked |= x == b; });
 		adj[a].push_back(b);
 		adj[b].push_back(a);
+		return !linked;
 	}
 
 	int size(int a) {
@@ -44,22 +44,38 @@ struct Naive {
 		dfs(a, [&](int /**/){res++;});
 		return res;
 	}
+
+	int add() {
+		int idx = ssize(adj);
+		adj.emplace_back();
+		seen.push_back(counter);
+		return idx;
+	}
 };
 
 void stress_test() {
 	ll queries = 0;
 	for (int tries = 0; tries < 200; tries++) {
 		int n = Random::integer<int>(1, 100);
-		UF uf(n);
+		UnionFind uf(n);
 		Naive naive(n);
-		for (int i = 0; i < n; i++) {
+		int rounds = n;
+		for (int i = 0; i < rounds; i++) {
 			for (int j = 0; j < 10; j++) {
 				int a = Random::integer<int>(0, n);
 				int b = Random::integer<int>(0, n);
-				uf.unionSets(a, b);
-				naive.unionSets(a, b);
+				auto got = uf.link(a, b);
+				auto expected = naive.link(a, b);
+				if (got != expected) cerr << "got: " << got << ", expected: " << expected << FAIL;
 			}
-			UF tmp = uf;
+			{
+				auto got = uf.add();
+				auto expected = naive.add();
+				assert(expected == n);
+				if (got != expected) cerr << "got: " << got << ", expected: " << expected << FAIL;
+				n++;
+			}
+			UnionFind tmp = uf;
 			for (int j = 0; j < n; j++) {
 				{
 					auto got = tmp.size(j);
@@ -69,8 +85,8 @@ void stress_test() {
 				{
 					int a = Random::integer<int>(0, n);
 					int b = Random::integer<int>(0, n);
-					bool got = tmp.findSet(a) == tmp.findSet(b);
-					bool expected = naive.findSet(a) == naive.findSet(b);
+					bool got = tmp.find(a) == tmp.find(b);
+					bool expected = naive.find(a) == naive.find(b);
 					if (got != expected) cerr << "got: " << got << ", expected: " << expected << FAIL;
 				}
 			}
@@ -84,7 +100,7 @@ constexpr int N = 2'000'000;
 void performance_test() {
 	timer t;
 	t.start();
-	UF uf(N);
+	UnionFind uf(N);
 	t.stop();
 	hash_t hash = 0;
 	for (int operations = 0; operations < N; operations++) {
@@ -92,9 +108,9 @@ void performance_test() {
 		int j = Random::integer<int>(0, N);
 		int k = Random::integer<int>(0, N);
 		int l = Random::integer<int>(0, N);
-		
+
 		t.start();
-		uf.unionSets(i, j);
+		uf.link(i, j);
 		hash += uf.size(k);
 		hash += uf.size(l);
 		t.stop();
diff --git a/test/graph/articulationPoints.bcc.cpp b/test/graph/articulationPoints.bcc.cpp
index cee2d0b..f112338 100644
--- a/test/graph/articulationPoints.bcc.cpp
+++ b/test/graph/articulationPoints.bcc.cpp
@@ -8,7 +8,7 @@ struct edge {
 #include <datastructures/unionFind.cpp>
 
 vector<vector<int>> naiveBCC(int m) {
-	init(m);
+	UnionFind uf(m);
 
 	vector<int> seen(ssize(adj), -1);
 	int run = 0;
@@ -28,13 +28,13 @@ vector<vector<int>> naiveBCC(int m) {
 				}
 			}
 			for (auto ee : adj[i]) {
-				if (seen[ee.to] == run) unionSets(ee.id, e.id);
+				if (seen[ee.to] == run) uf.link(ee.id, e.id);
 			}
 		}
 	}
 	vector<vector<int>> res(m);
 	for (int i = 0; i < m; i++) {
-		res[findSet(i)].push_back(i);
+		res[uf.find(i)].push_back(i);
 	}
 	for (auto& v : res) ranges::sort(v);
 	res.erase(begin(ranges::remove_if(res, [](const vector<int>& v){return ssize(v) <= 1;})), end(res));
diff --git a/test/graph/cycleCounting.cpp b/test/graph/cycleCounting.cpp
index 9c7bf0c..82caf16 100644
--- a/test/graph/cycleCounting.cpp
+++ b/test/graph/cycleCounting.cpp
@@ -6,18 +6,14 @@ int naive(const vector<pair<int, int>>& edges, int n) {
 	int res = 0;
 	for (int i = 1; i < (1ll << ssize(edges)); i++) {
 		vector<int> deg(n);
-		init(n);
+		UnionFind uf(n);
 		int cycles = 0;
 		for (int j = 0; j < ssize(edges); j++) {
 			if (((i >> j) & 1) != 0) {
 				auto [a, b] = edges[j];
 				deg[a]++;
 				deg[b]++;
-				if (findSet(a) != findSet(b)) {
-					unionSets(a, b);
-				} else {
-					cycles++;
-				}
+				if (!uf.link(a, b)) cycles++;
 			}
 		}
 		bool ok = cycles == 1;
diff --git a/test/graph/kruskal.cpp b/test/graph/kruskal.cpp
index f6245b9..bc1cce5 100644
--- a/test/graph/kruskal.cpp
+++ b/test/graph/kruskal.cpp
@@ -1,22 +1,19 @@
 #include "../util.h"
 #include <datastructures/unionFind.cpp>
 
-struct edge {
+#define Edge Edge_ // we have a struct named Edge in util.h
+
+struct Edge {
 	int from, to;
 	ll cost;
-	bool operator<(const edge& o) const {
+	bool operator<(const Edge& o) const {
 		return cost > o.cost;
 	}
 };
-ll kruskal(vector<edge>& edges, int n) {
-	init(n);
-	#define Edge edge
-	#include <graph/kruskal.cpp>
-	#undef Edge
-	return cost;
-}
 
-ll prim(vector<edge>& edges, int n) {
+#include <graph/kruskal.cpp>
+
+ll prim(vector<Edge>& edges, int n) {
 	vector<vector<pair<ll, int>>> adj(n);
 	for (auto [a, b, d] : edges) {
 		adj[a].emplace_back(d, b);
@@ -51,13 +48,14 @@ void stress_test() {
 
 		Graph<NoData> g(n);
 		g.erdosRenyi(m);
-		vector<edge> edges;
+		vector<Edge> edges;
 		g.forEdges([&](int a, int b){
 			ll w = Random::integer<ll>(-1'000'000'000ll, 1'000'000'000ll);
 			edges.push_back({a, b, w});
 		});
 
-		ll got = kruskal(edges, n);
+		vector<Edge> mst;
+		ll got = kruskal(n, edges, mst);
 		ll expected = prim(edges, n);
 
 		if (got != expected) cerr << "got: " << got << ", expected: " << expected << FAIL;
@@ -72,14 +70,15 @@ void performance_test() {
 	timer t;
 	Graph<NoData> g(N);
 	g.erdosRenyi(M);
-	vector<edge> edges;
+	vector<Edge> edges;
 	g.forEdges([&](int a, int b){
 		ll w = Random::integer<ll>(-1'000'000'000ll, 1'000'000'000ll);
 		edges.push_back({a, b, w});
 	});
 
 	t.start();
-	hash_t hash = kruskal(edges, N);
+	vector<Edge> mst;
+	hash_t hash = kruskal(N, edges, mst);
 	t.stop();
 	if (t.time > 1000) cerr << "too slow: " << t.time << FAIL;
 	cerr << "tested performance: " << t.time << "ms (hash: " << hash << ")" << endl;