Persistent DSU (Fully Persistent Union-Find)

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
// Persistent segment tree storing pairs {parent, rank}
5
struct Node {
6
    Node *l, *r;
7
    int parent;
8
    int rnk; // rank (or size-log)
9
    Node(int p=0, int rk=0) : l(nullptr), r(nullptr), parent(p), rnk(rk) {}
10
};
11
12
struct PersistentArray {
13
    int n;
14
    vector<Node*> roots; // roots[version]
15
16
    PersistentArray(int n): n(n) {
17
        roots.clear();
18
        roots.push_back(build(1, n)); // version 0
19
    }
20
21
    Node* build(int L, int R) {
22
        Node* node = new Node();
23
        if (L == R) {
24
            node->parent = L; // initially parent[i] = i
25
            node->rnk = 0;    // rank[i] = 0
26
            return node;
27
        }
28
        int M = (L + R) >> 1;
29
        node->l = build(L, M);
30
        node->r = build(M+1, R);
31
        return node;
32
    }
33
34
    // Read at index idx from version v
35
    pair<int,int> get(Node* root, int L, int R, int idx) const {
36
        if (L == R) return {root->parent, root->rnk};
37
        int M = (L + R) >> 1;
38
        if (idx <= M) return get(root->l, L, M, idx);
39
        else return get(root->r, M+1, R, idx);
40
    }
41
42
    pair<int,int> get(int version, int idx) const {
43
        return get(roots[version], 1, n, idx);
44
    }
45
46
    // Point-update at idx: set parent or rank or both; returns new root
47
    Node* set(Node* root, int L, int R, int idx, int newParent, int newRank) {
48
        Node* node = new Node();
49
        if (L == R) {
50
            node->parent = newParent;
51
            node->rnk = newRank;
52
            return node;
53
        }
54
        int M = (L + R) >> 1;
55
        if (idx <= M) {
56
            node->l = set(root->l, L, M, idx, newParent, newRank);
57
            node->r = root->r;
58
        } else {
59
            node->l = root->l;
60
            node->r = set(root->r, M+1, R, idx, newParent, newRank);
61
        }
62
        return node;
63
    }
64
65
    // Convenience: write only parent or only rank by first reading
66
    Node* write_parent(Node* root, int idx, int L, int R, int newParent) {
67
        auto cur = get(root, L, R, idx);
68
        return set(root, L, R, idx, newParent, cur.second);
69
    }
70
    Node* write_rank(Node* root, int idx, int L, int R, int newRank) {
71
        auto cur = get(root, L, R, idx);
72
        return set(root, L, R, idx, cur.first, newRank);
73
    }
74
75
    // Find representative of x in given root version; no path compression
76
    int find_rep(Node* root, int x) const {
77
        while (true) {
78
            auto pr = get(root, 1, n, x);
79
            int p = pr.first;
80
            if (p == x) return x;
81
            x = p;
82
        }
83
    }
84
85
    // Union by rank on baseVersion; returns index of new version
86
    int unite(int baseVersion, int a, int b) {
87
        Node* baseRoot = roots[baseVersion];
88
        int ra = find_rep(baseRoot, a);
89
        int rb = find_rep(baseRoot, b);
90
        if (ra == rb) {
91
            // No change; just reuse the base root as a new version (branching)
92
            roots.push_back(baseRoot);
93
            return (int)roots.size() - 1;
94
        }
95
        auto pa = get(baseRoot, 1, n, ra);
96
        auto pb = get(baseRoot, 1, n, rb);
97
        int rka = pa.second;
98
        int rkb = pb.second;
99
        Node* newRoot = baseRoot;
100
        if (rka < rkb) {
101
            // Attach ra under rb
102
            newRoot = write_parent(newRoot, ra, 1, n, rb);
103
        } else if (rka > rkb) {
104
            // Attach rb under ra
105
            newRoot = write_parent(newRoot, rb, 1, n, ra);
106
        } else {
107
            // Equal rank: attach rb under ra, increment rank[ra]
108
            newRoot = write_parent(newRoot, rb, 1, n, ra);
109
            newRoot = write_rank(newRoot, ra, 1, n, rka + 1);
110
        }
111
        roots.push_back(newRoot);
112
        return (int)roots.size() - 1;
113
    }
114
115
    // Connectivity query on a given version
116
    bool connected(int version, int a, int b) const {
117
        Node* root = roots[version];
118
        return find_rep(root, a) == find_rep(root, b);
119
    }
120
};
121
122
int main() {
123
    ios::sync_with_stdio(false);
124
    cin.tie(nullptr);
125
126
    // Demo: n elements, build version 0; then create branches
127
    int n = 8;
128
    PersistentArray dsu(n);
129
130
    // Version timeline:
131
    // v0: initial
132
    int v1 = dsu.unite(0, 1, 2);            // v1 = v0 + union(1,2)
133
    int v2 = dsu.unite(v1, 2, 3);           // v2 = v1 + union(2,3)
134
    int v3 = dsu.unite(0, 5, 6);            // branch from v0: v3 = v0 + union(5,6)
135
    int v4 = dsu.unite(v2, 1, 3);           // union on already-connected nodes; version still created
136
137
    cout << boolalpha;
138
    cout << "v0: conn(1,3) = " << dsu.connected(0, 1, 3) << "\n"; // false
139
    cout << "v1: conn(1,3) = " << dsu.connected(v1, 1, 3) << "\n"; // false
140
    cout << "v2: conn(1,3) = " << dsu.connected(v2, 1, 3) << "\n"; // true
141
    cout << "v3: conn(5,6) = " << dsu.connected(v3, 5, 6) << "\n"; // true
142
    cout << "v4: conn(1,3) = " << dsu.connected(v4, 1, 3) << "\n"; // true
143
144
    // You can continue branching from any version index stored in dsu.roots
145
    return 0;
146
}
147

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
struct DSURollback {
5
    int n;
6
    vector<int> parent, sz;
7
    // Stack of changes: (v, old_parent, u, old_size). We use a small struct.
8
    struct Change { int v, old_parent; int u, old_size; bool merged; };
9
    vector<Change> st;
10
11
    DSURollback(int n=0): n(n), parent(n+1), sz(n+1) {
12
        iota(parent.begin(), parent.end(), 0);
13
        fill(sz.begin(), sz.end(), 1);
14
    }
15
16
    int find(int x) {
17
        while (x != parent[x]) x = parent[x];
18
        return x;
19
    }
20
21
    int snapshot() const { return (int)st.size(); }
22
23
    void rollback(int snap) {
24
        while ((int)st.size() > snap) {
25
            auto c = st.back(); st.pop_back();
26
            if (!c.merged) continue; // no-op union
27
            parent[c.v] = c.old_parent;
28
            sz[c.u] = c.old_size;
29
        }
30
    }
31
32
    void unite(int a, int b) {
33
        a = find(a); b = find(b);
34
        if (a == b) { st.push_back({0,0,0,0,false}); return; }
35
        if (sz[a] < sz[b]) swap(a,b);
36
        // attach b under a
37
        st.push_back({b, parent[b], a, sz[a], true});
38
        parent[b] = a;
39
        sz[a] += sz[b];
40
    }
41
};
42
43
// Offline dynamic connectivity: supports add/remove edges and connectivity queries
44
// Input format (example in comments):
45
// Q queries. Types:
46
//  + u v  : add edge (u,v)
47
//  - u v  : remove edge (u,v)
48
//  ? u v  : query connectivity of u and v at current time
49
// Assumption: Each edge added will be removed at most once (standard), unmatched adds persist to the end.
50
51
struct OfflineDC {
52
    int n, Q;
53
    vector<vector<pair<int,int>>> seg; // segment tree over time storing edges
54
    vector<pair<char, pair<int,int>>> ops; // queries in order
55
    map<pair<int,int>, int> open; // edge -> start time
56
    vector<string> answers;
57
58
    OfflineDC(int n, int Q): n(n), Q(Q) {
59
        seg.assign(4*Q+5, {});
60
        ops.reserve(Q+1);
61
    }
62
63
    static pair<int,int> norm(int u, int v) { if (u>v) swap(u,v); return {u,v}; }
64
65
    void add_op(char t, int u, int v) { ops.push_back({t, {u,v}}); }
66
67
    void add_interval(int idx, int L, int R, int ql, int qr, pair<int,int> e) {
68
        if (ql>R || qr<L) return;
69
        if (ql<=L && R<=qr) { seg[idx].push_back(e); return; }
70
        int M=(L+R)>>1;
71
        add_interval(idx<<1, L, M, ql, qr, e);
72
        add_interval(idx<<1|1, M+1, R, ql, qr, e);
73
    }
74
75
    void build_intervals() {
76
        for (int t=1; t<=Q; ++t) {
77
            char c = ops[t-1].first;
78
            int u = ops[t-1].second.first;
79
            int v = ops[t-1].second.second;
80
            if (c == '+') {
81
                open[norm(u,v)] = t;
82
            } else if (c == '-') {
83
                auto e = norm(u,v);
84
                auto it = open.find(e);
85
                if (it != open.end()) {
86
                    int s = it->second;
87
                    add_interval(1, 1, Q, s, t-1, e);
88
                    open.erase(it);
89
                }
90
            }
91
        }
92
        for (auto &kv : open) {
93
            add_interval(1, 1, Q, kv.second, Q, kv.first);
94
        }
95
        open.clear();
96
    }
97
98
    void dfs(int idx, int L, int R, DSURollback &dsu) {
99
        int snap = dsu.snapshot();
100
        for (auto &e : seg[idx]) dsu.unite(e.first, e.second);
101
        if (L == R) {
102
            char c = ops[L-1].first;
103
            int u = ops[L-1].second.first;
104
            int v = ops[L-1].second.second;
105
            if (c == '?') {
106
                bool ans = (dsu.find(u) == dsu.find(v));
107
                answers.push_back(ans ? "YES" : "NO");
108
            }
109
            dsu.rollback(snap);
110
            return;
111
        }
112
        int M=(L+R)>>1;
113
        dfs(idx<<1, L, M, dsu);
114
        dfs(idx<<1|1, M+1, R, dsu);
115
        dsu.rollback(snap);
116
    }
117
118
    vector<string> solve() {
119
        build_intervals();
120
        DSURollback dsu(n);
121
        answers.clear();
122
        dfs(1, 1, Q, dsu);
123
        return answers;
124
    }
125
};
126
127
int main(){
128
    ios::sync_with_stdio(false);
129
    cin.tie(nullptr);
130
131
    // Example usage:
132
    // n=5, Q=7
133
    // + 1 2
134
    // + 2 3
135
    // ? 1 3   -> YES
136
    // - 2 3
137
    // ? 1 3   -> NO
138
    // + 4 5
139
    // ? 4 5   -> YES
140
141
    int n = 5, Q = 7;
142
    OfflineDC solver(n, Q);
143
    vector<tuple<char,int,int>> seq = {
144
        {'+',1,2}, {'+',2,3}, {'?',1,3}, {'-',2,3}, {'?',1,3}, {'+',4,5}, {'?',4,5}
145
    };
146
    for (auto [c,u,v] : seq) solver.add_op(c,u,v);
147
148
    auto ans = solver.solve();
149
    for (auto &s : ans) cout << s << '\n';
150
    return 0;
151
}
152