unionfind/fully_persistent_uf.hpp

• View this file on GitHub
• Last update: 2022-01-08 20:23:44+09:00
• Include: #include "unionfind/fully_persistent_uf.hpp"
• Link: https://ei1333.github.io/library/structure/union-find/persistent-union-find.cpp>

Depends on

• data_structure/persistent_array.hpp

Verified with

• unionfind/test/fully_persistent_uf.test.cpp

Code

#pragma once
#include "../data_structure/persistent_array.hpp"
#include <vector>

// CUT begin
// Fully persistent UnionFind
// - find(t, x) : find the root of DSU tree x belongs to, when unite() is called t times.
// Complexity: O(logN) for each query
// Reference: <https://ei1333.github.io/library/structure/union-find/persistent-union-find.cpp>
struct PersistentUnionFind {
    int N;
    using Array = persistent_array<int, 4>;
    Array par;
    std::vector<Array::node *> savepoints; // Tree structure is saved after every `unite()` operation
    PersistentUnionFind(int N) : N(N), par(-1) { savepoints.emplace_back(nullptr); }
    int find(int t, int x) const {
        const int p = par.get(savepoints[t], x);
        return p < 0 ? x : find(t, p);
    }
    bool unite(int t, int x, int y) {
        x = find(t, x), y = find(t, y);
        auto ptr = savepoints[t];
        if (x != y) {
            int sz_x = -par.get(savepoints[t], x), sz_y = -par.get(savepoints[t], y);
            if (sz_x > sz_y) {
                par.ch(ptr, x, -sz_x - sz_y), par.ch(ptr, y, x);
            } else {
                par.ch(ptr, y, -sz_x - sz_y), par.ch(ptr, x, y);
            }
        }
        return savepoints.emplace_back(ptr), x != y;
    }
    int count(int t, int x) const { return -par.get(savepoints[t], find(t, x)); }
    bool same(int t, int x, int y) const { return find(t, x) == find(t, y); }
};

#line 2 "data_structure/persistent_array.hpp"
#include <algorithm>
#include <array>
#include <vector>

// CUT begin
// Persistent Array
// Reference: <https://qiita.com/hotman78/items/9c643feae1de087e6fc5>
//            <https://ei1333.github.io/luzhiled/snippets/structure/persistent-array.html>
// - (2^LOG)-ary tree-based
// - Fully persistent
// - `get(root, k)`:  get k-th element
// - `set(root, k, data)`: make new array whose k-th element is updated to data
// - `ch(root, k, data)` : change k-th element and implicitly update root
// Verify: Codeforces 707D <https://codeforces.com/contest/707/problem/D>
template <typename T, int LOG> struct persistent_array {
    T zero;
    struct node {
        T data;
        std::array<node *, 1 << LOG> child;
        node(const T &d) : data(d) { std::fill(child.begin(), child.end(), nullptr); }
        node(const node &n) : data(n.data) {
            std::copy(n.child.begin(), n.child.end(), child.begin());
        }
    };
    persistent_array(T zero) : zero(zero) {}

    T get(node *t, int k) const {
        if (t == nullptr) { return zero; }
        return k ? get(t->child[k & ((1 << LOG) - 1)], k >> LOG) : t->data;
    }

    [[nodiscard]] node *set(node *t, int k, const T &data) {
        t = (t == nullptr) ? new node(zero) : new node(*t);
        if (k == 0) {
            t->data = data;
        } else {
            auto ptr = set(t->child[k & ((1 << LOG) - 1)], k >> LOG, data);
            t->child[k & ((1 << LOG) - 1)] = ptr;
        }
        return t;
    }

    void ch(node *&t, int k, const T &data) { t = set(t, k, data); }

    [[nodiscard]] node *build(const std::vector<T> &vec) {
        node *root = nullptr;
        for (size_t i = 0; i < vec.size(); i++) { root = set(root, i, vec[i]); }
        return root;
    }
};
#line 4 "unionfind/fully_persistent_uf.hpp"

// CUT begin
// Fully persistent UnionFind
// - find(t, x) : find the root of DSU tree x belongs to, when unite() is called t times.
// Complexity: O(logN) for each query
// Reference: <https://ei1333.github.io/library/structure/union-find/persistent-union-find.cpp>
struct PersistentUnionFind {
    int N;
    using Array = persistent_array<int, 4>;
    Array par;
    std::vector<Array::node *> savepoints; // Tree structure is saved after every `unite()` operation
    PersistentUnionFind(int N) : N(N), par(-1) { savepoints.emplace_back(nullptr); }
    int find(int t, int x) const {
        const int p = par.get(savepoints[t], x);
        return p < 0 ? x : find(t, p);
    }
    bool unite(int t, int x, int y) {
        x = find(t, x), y = find(t, y);
        auto ptr = savepoints[t];
        if (x != y) {
            int sz_x = -par.get(savepoints[t], x), sz_y = -par.get(savepoints[t], y);
            if (sz_x > sz_y) {
                par.ch(ptr, x, -sz_x - sz_y), par.ch(ptr, y, x);
            } else {
                par.ch(ptr, y, -sz_x - sz_y), par.ch(ptr, x, y);
            }
        }
        return savepoints.emplace_back(ptr), x != y;
    }
    int count(int t, int x) const { return -par.get(savepoints[t], find(t, x)); }
    bool same(int t, int x, int y) const { return find(t, x) == find(t, y); }
};

Back to top page