test/yosupo/persistent_unionfind.test.cpp

View this file on GitHub
Last update: 2024-01-07 20:49:49+09:00
Problem: https://judge.yosupo.jp/problem/persistent_unionfind

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/persistent_unionfind"

#include "../../data-structure/unionfind/persistent_union_find.hpp"

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    cin >> N >> Q;
    PersistentUnionFind init(N);
    vector<PersistentUnionFind> ver(Q);
    for (int i = 0; i < Q; ++i) {
        int t, k, u, v;
        cin >> t >> k >> u >> v;
        if (t == 0) {
            if (k == -1) ver[i] = init.unite(u, v);
            else ver[i] = ver[k].unite(u, v);
        } else {
            int b;
            if (k == -1) b = init.same(u, v);
            else b = ver[k].same(u, v);
            cout << b << "\n";
        }
    }
}

#line 1 "test/yosupo/persistent_unionfind.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/persistent_unionfind"

#line 2 "data-structure/unionfind/persistent_union_find.hpp"
#include <algorithm>
#include <vector>

#line 2 "data-structure/persistent_array.hpp"
#include <memory>
#line 4 "data-structure/persistent_array.hpp"

template <typename T, int B = 2>
class PersistentArray {
   public:
    PersistentArray() = default;
    explicit PersistentArray(const std::vector<T>& v) {
        for (int i = 0; i < (int)v.size(); ++i) root = set(root, i, v[i]);
    }

    T get(int k) const { return get(root, k); }

    PersistentArray set(int k, const T& x) const {
        return PersistentArray(set(root, k, x));
    }

   private:
    struct Node;
    using node_ptr = std::shared_ptr<Node>;

    struct Node {
        T val;
        node_ptr ch[B];
    };

    node_ptr root = nullptr;

    explicit PersistentArray(const node_ptr& root) : root(root) {}

    T get(const node_ptr& t, int k) const {
        if (k == 0) return t->val;
        return get(t->ch[k % B], k / B);
    }

    node_ptr set(const node_ptr& t, int k, const T& x) const {
        node_ptr res =
            t ? std::make_shared<Node>(*t) : std::make_shared<Node>();
        if (k == 0) {
            res->val = x;
        } else {
            res->ch[k % B] = set(res->ch[k % B], k / B, x);
        }
        return res;
    }
};
#line 6 "data-structure/unionfind/persistent_union_find.hpp"

class PersistentUnionFind {
   public:
    PersistentUnionFind() = default;
    explicit PersistentUnionFind(int n) : data(std::vector<int>(n, -1)) {}

    int find(int x) const {
        int p = data.get(x);
        if (p < 0) return x;
        return find(p);
    }

    PersistentUnionFind unite(int x, int y) const {
        x = find(x);
        y = find(y);
        if (x == y) return *this;
        int u = data.get(x);
        int v = data.get(y);
        if (u > v) std::swap(x, y);
        auto tmp = data.set(x, u + v);
        auto res = tmp.set(y, x);
        return PersistentUnionFind(res);
    }

    bool same(int x, int y) const { return find(x) == find(y); }

    int size(int x) const { return -data.get(x); }

   private:
    PersistentArray<int, 8> data;

    explicit PersistentUnionFind(const PersistentArray<int, 8>& pa)
        : data(pa) {}
};
#line 4 "test/yosupo/persistent_unionfind.test.cpp"

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    cin >> N >> Q;
    PersistentUnionFind init(N);
    vector<PersistentUnionFind> ver(Q);
    for (int i = 0; i < Q; ++i) {
        int t, k, u, v;
        cin >> t >> k >> u >> v;
        if (t == 0) {
            if (k == -1) ver[i] = init.unite(u, v);
            else ver[i] = ver[k].unite(u, v);
        } else {
            int b;
            if (k == -1) b = init.same(u, v);
            else b = ver[k].same(u, v);
            cout << b << "\n";
        }
    }
}