test/rollback_union_find.test.cpp

View this file on GitHub
Last update: 2021-12-28 22:38:25+09:00
Problem: https://judge.yosupo.jp/problem/persistent_unionfind

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/persistent_unionfind"
#include "../utility/rep.cpp"
#include "../utility/rec_lambda.cpp"
#include "../graph/rollback_union_find.cpp"
#include <iostream>

int main() {
    int N, Q;
    std::cin >> N >> Q;
    std::vector<std::pair<int, int>> connect(Q + 1);
    std::vector<std::vector<int>> graph(Q + 1);
    std::vector<std::vector<std::tuple<int, int, int>>> query(Q + 1);
    std::vector<bool> ans(Q + 1);
    std::vector<int> output;
    output.reserve(Q);
    for (const auto i: rep(1, Q + 1)) {
        int t, k;
        int u, v;
        std::cin >> t >> k >> u >> v;
        k += 1;
        if (t == 0) {
            connect[i] = {u, v};
            graph[k].push_back(i);
        }
        else {
            output.push_back(i);
            query[k].emplace_back(i, u, v);
        }
    }
    RollbackUnionFind dsu(N);
    std::vector<bool> done(Q);
    rec_lambda([&](auto&& dfs, const int u) -> void {
        const auto f = dsu.merge(connect[u].first, connect[u].second).second;
        for (const auto [i, x, y]: query[u]) {
            ans[i] = dsu.same(x, y);
        }
        for (const auto v: graph[u]) {
            dfs(v);
        }
        dsu.rollback(f);
        done[u] = true;
    })(0);
    for (const auto i: output) {
        std::cout << ans[i] << '\n';
    }
}

#line 1 "test/rollback_union_find.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/persistent_unionfind"
#line 2 "utility/rep.cpp"
#include <algorithm>

class Range {
    struct Iter {
        int itr;
        constexpr Iter(const int pos) noexcept : itr(pos) {}
        constexpr void operator++() noexcept { ++itr; }
        constexpr bool operator!=(const Iter& other) const noexcept { return itr != other.itr; }
        constexpr int operator*() const noexcept { return itr; }
    };
    const Iter first, last;

  public:
    explicit constexpr Range(const int first, const int last) noexcept : first(first), last(std::max(first, last)) {}
    constexpr Iter begin() const noexcept { return first; }
    constexpr Iter end() const noexcept { return last; }
};

constexpr Range rep(const int l, const int r) noexcept { return Range(l, r); }
constexpr Range rep(const int n) noexcept { return Range(0, n); }
#line 2 "utility/rec_lambda.cpp"
#include <type_traits>
#include <utility>

template <class F> struct RecursiveLambda : private F {
    explicit constexpr RecursiveLambda(F&& f) : F(std::forward<F>(f)) {}
    template <class... Args> constexpr decltype(auto) operator()(Args&&... args) const {
        return F::operator()(*this, std::forward<Args>(args)...);
    }
};

template <class F> constexpr decltype(auto) rec_lambda(F&& f) { return RecursiveLambda<F>(std::forward<F>(f)); }
#line 2 "graph/rollback_union_find.cpp"
#include <cassert>
#line 4 "graph/rollback_union_find.cpp"
#include <vector>

class RollbackUnionFind {
    std::vector<int> data;
    std::vector<std::pair<int, int>> history;

  public:
    explicit RollbackUnionFind(const int size = 0) : data(size, -1), history() {}

    int size() const { return data.size(); }

    int leader(int u) const {
        assert(0 <= u and u < size());
        while (data[u] >= 0) u = data[u];
        return u;
    }

    int size(const int u) const {
        assert(0 <= u and u < size());
        return -data[leader(u)];
    }

    std::pair<int, bool> merge(int u, int v) {
        assert(0 <= u and u < size());
        assert(0 <= v and v < size());
        u = leader(u);
        v = leader(v);
        if (u == v) return std::make_pair(u, false);
        if (data[u] > data[v]) std::swap(u, v);
        history.emplace_back(u, data[u]);
        history.emplace_back(v, data[v]);
        data[u] += data[v];
        data[v] = u;
        return std::make_pair(u, true);
    }

    bool same(const int u, const int v) const {
        assert(0 <= u and u < size());
        assert(0 <= v and v < size());
        return leader(u) == leader(v);
    }

    void rollback(const int steps) {
        assert(0 <= steps and 2 * steps <= (int)history.size());
        for (int i = 2 * steps; i > 0; --i) {
            const auto [k, x] = history.back();
            history.pop_back();
            data[k] = x;
        }
    }
};
#line 5 "test/rollback_union_find.test.cpp"
#include <iostream>

int main() {
    int N, Q;
    std::cin >> N >> Q;
    std::vector<std::pair<int, int>> connect(Q + 1);
    std::vector<std::vector<int>> graph(Q + 1);
    std::vector<std::vector<std::tuple<int, int, int>>> query(Q + 1);
    std::vector<bool> ans(Q + 1);
    std::vector<int> output;
    output.reserve(Q);
    for (const auto i: rep(1, Q + 1)) {
        int t, k;
        int u, v;
        std::cin >> t >> k >> u >> v;
        k += 1;
        if (t == 0) {
            connect[i] = {u, v};
            graph[k].push_back(i);
        }
        else {
            output.push_back(i);
            query[k].emplace_back(i, u, v);
        }
    }
    RollbackUnionFind dsu(N);
    std::vector<bool> done(Q);
    rec_lambda([&](auto&& dfs, const int u) -> void {
        const auto f = dsu.merge(connect[u].first, connect[u].second).second;
        for (const auto [i, x, y]: query[u]) {
            ans[i] = dsu.same(x, y);
        }
        for (const auto v: graph[u]) {
            dfs(v);
        }
        dsu.rollback(f);
        done[u] = true;
    })(0);
    for (const auto i: output) {
        std::cout << ans[i] << '\n';
    }
}