DataStructure/DSU/DSU_persistent.h

• View this file on GitHub
• Last update: 2023-01-07 01:46:12+08:00
• Link: https://judge.yosupo.jp/problem/persistent_unionfind
• Link: https://oj.vnoi.info/problem/hello22_schoolplan)

Depends on

• DataStructure/PersistentArray.h

Verified with

• DataStructure/test/persistent_dsu.test.cpp

Code

// PersistentDSU
//
// Notes:
// - this doesn't support delete edge operation, so isn't enough to
//   solve dynamic connectivity problem.
// - it has high mem and time usage, so be careful (both TLE and MLE on
//   https://oj.vnoi.info/problem/hello22_schoolplan)
//
// Tested:
// - https://judge.yosupo.jp/problem/persistent_unionfind
#include "../PersistentArray.h"
struct PersistentDSU {
    int n;
    using Arr = PersistentArray<int>;

    PersistentDSU(int _n) : n(_n) {
        roots.emplace_back(A.build(std::vector<int> (n, -1)));
    }

    int find(int version, int u) {
        // Note that we can't do path compression here
        int p = A.get(roots[version], u);
        return p < 0 ? u : find(version, p);
    }

    // Note that this will always create a new version,
    // regardless of whether u and v was previously in same component.
    bool merge(int version, int u, int v) {
        u = find(version, u);
        v = find(version, v);
        auto ptr = roots[version];
        if (u != v) {
            int sz_u = -A.get(ptr, u), sz_v = -A.get(ptr, v);
            if (sz_u < sz_v) swap(u, v);
            // sz[u] >= sz[v]
            ptr = A.set(ptr, u, -sz_u - sz_v);
            ptr = A.set(ptr, v, u);
        }

        roots.emplace_back(ptr);
        return u != v;
    }

    int component_size(int version, int u) {
        return -A.get(roots[version], find(version, u));
    }

    bool same_component(int version, int u, int v) {
        return find(version, u) == find(version, v);
    }

    Arr A;
    vector<Arr::Node*> roots;
};

#line 1 "DataStructure/DSU/DSU_persistent.h"
// PersistentDSU
//
// Notes:
// - this doesn't support delete edge operation, so isn't enough to
//   solve dynamic connectivity problem.
// - it has high mem and time usage, so be careful (both TLE and MLE on
//   https://oj.vnoi.info/problem/hello22_schoolplan)
//
// Tested:
// - https://judge.yosupo.jp/problem/persistent_unionfind
#line 1 "DataStructure/PersistentArray.h"
// PersistentArray
//
// Notes:
// - Reduce mem -> decrease LOG
// - Reduce time -> increase LOG
//
// Tested:
// - https://codeforces.com/contest/707/problem/D
template<typename T>
struct PersistentArray {
    static const int LOG = 4;
    static const int FULL_MASK = (1<<LOG) - 1;

    struct Node {
        T x;
        std::array<Node*, 1<<LOG> child;
        Node(const T& _x) : x(_x) {
            std::fill(child.begin(), child.end(), nullptr);
        }
        Node(const Node& n) : x(n.x) {
            std::copy(n.child.begin(), n.child.end(), child.begin());
        }
    };

    // get p-th element in `t` version
    static T get(Node* t, int p) {
        if (t == NULL) return 0;
        return p ? get(t->child[p & FULL_MASK], p >> LOG) : t->x;
    }

    // set p-th element in `t` version to x, and return new version
    static Node* set(Node* t, int p, const T& x) {
        t = (t == NULL) ? new Node(0) : new Node(*t);
        if (p == 0) t->x = x;
        else {
            auto ptr = set(t->child[p & FULL_MASK], p >> LOG, x);
            t->child[p & FULL_MASK] = ptr;
        }
        return t;
    }

    // init a persistent array and return root node
    Node* build(const vector<T>& v) {
        Node* root = NULL;
        for (int i = 0; i < (int) v.size(); i++) {
            root = set(root, i, v[i]);
        }
        return root;
    }
};
#line 12 "DataStructure/DSU/DSU_persistent.h"
struct PersistentDSU {
    int n;
    using Arr = PersistentArray<int>;

    PersistentDSU(int _n) : n(_n) {
        roots.emplace_back(A.build(std::vector<int> (n, -1)));
    }

    int find(int version, int u) {
        // Note that we can't do path compression here
        int p = A.get(roots[version], u);
        return p < 0 ? u : find(version, p);
    }

    // Note that this will always create a new version,
    // regardless of whether u and v was previously in same component.
    bool merge(int version, int u, int v) {
        u = find(version, u);
        v = find(version, v);
        auto ptr = roots[version];
        if (u != v) {
            int sz_u = -A.get(ptr, u), sz_v = -A.get(ptr, v);
            if (sz_u < sz_v) swap(u, v);
            // sz[u] >= sz[v]
            ptr = A.set(ptr, u, -sz_u - sz_v);
            ptr = A.set(ptr, v, u);
        }

        roots.emplace_back(ptr);
        return u != v;
    }

    int component_size(int version, int u) {
        return -A.get(roots[version], find(version, u));
    }

    bool same_component(int version, int u, int v) {
        return find(version, u) == find(version, v);
    }

    Arr A;
    vector<Arr::Node*> roots;
};

Back to top page