From 35a2542e4876979c9d5275f360c59b6d85c1a616 Mon Sep 17 00:00:00 2001
From: subcrip <contact@subc.rip>
Date: Sat, 5 Oct 2024 19:25:39 +0800
Subject: [PATCH] Add trees/quick_union_rollback.cc

Signed-off-by: subcrip <contact@subc.rip>
---
 trees/quick_union_rollback.cc | 43 +++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 trees/quick_union_rollback.cc

diff --git a/trees/quick_union_rollback.cc b/trees/quick_union_rollback.cc
new file mode 100644
index 0000000..07c190b
--- /dev/null
+++ b/trees/quick_union_rollback.cc
@@ -0,0 +1,43 @@
+class quick_union {
+public:
+    size_t n;
+    vector<size_t> c, sz;
+    vector<pair<size_t, size_t>> history;
+public:
+    quick_union(size_t n) : n(n), c(n), sz(n) {
+        iota(c.begin(), c.end(), 0);
+        sz.assign(n, 1);
+    }
+
+    size_t query(size_t i) {
+        if (c[i] != i) return query(c[i]);
+        return c[i];
+    }
+
+    void merge(size_t i, size_t j) {
+        if (connected(i, j)) {
+            history.emplace_back(n, n);
+        } else {
+            i = query(i), j = query(j);
+            if (sz[i] > sz[j]) swap(i, j);
+            history.emplace_back(i, j);
+            sz[j] += sz[i];
+            c[i] = j;
+        }
+    }
+
+    bool connected(size_t i, size_t j) {
+        return query(i) == query(j);
+    }
+
+    size_t query_size(size_t i) {
+        return sz[query(i)];
+    }
+
+    void rollback() {
+        auto [i, j] = popback(history);
+        if (i == n and j == n) return;
+        c[i] = i;
+        sz[j] -= sz[i];
+    }
+};