UnionFindRollback

Disjoint-set data structure with undo. If undo is not needed, skip st, time() and rollback().

Time: O(log(N)) 21 lines tested as part of DirectedMST.h

Usage: int t = uf.time(); ...; uf.rollback(t);

content/data-structures/UnionFindRollback.h — Lukas Polacek, Simon Lindholm, source: folklore

struct RollbackUF {
	vi e; vector<pii> st;
	RollbackUF(int n) : e(n, -1) {}
	int size(int x) { return -e[find(x)]; }
	int find(int x) { return e[x] < 0 ? x : find(e[x]); }
	int time() { return sz(st); }
	void rollback(int t) {
		for (int i = time(); i --> t;)
			e[st[i].first] = st[i].second;
		st.resize(t);
	}
	bool join(int a, int b) {
		a = find(a), b = find(b);
		if (a == b) return false;
		if (e[a] > e[b]) swap(a, b);
		st.push_back({a, e[a]});
		st.push_back({b, e[b]});
		e[a] += e[b]; e[b] = a;
		return true;
	}
};