verify/AOJ-0723.test.cpp

• View this file on GitHub
• Last update: 2025-05-30 20:02:37+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/challenges/sources/JOI/Prelim/0723

Depends on

• Library/Common.hpp
• Rollback Union-Find - Rollback 可能 Union-Find (Library/DataStructure/RollbackUnionFind.hpp)

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/challenges/sources/JOI/Prelim/0723"

#include "../Library/DataStructure/RollbackUnionFind.hpp"

int main(){
    cin.tie(0)->sync_with_stdio(false);
    int N, M, K; cin >> N >> M >> K;
    vector<int> U(M), V(M);
    for(int i = 0; i < M; ++i) cin >> U[i] >> V[i], --U[i], --V[i];
    vector<int> S(N);
    for(int i = 0; i < N; ++i) cin >> S[i], --S[i];

    RollbackUnionFind uf(N);
    map<pair<int, int>, vector<pair<int, int>>> Road;
    for(int i = 0; i < M; ++i){
        int u = U[i], v = V[i];
        int su = S[u], sv = S[v];
        if(su == sv) uf.Unite(u, v);
        else{
            if(su > sv) swap(su, sv);
            Road[{su, sv}].push_back({u, v});
        }
    }
    int Q; cin >> Q;
    map<pair<int, int>, vector<tuple<int, int, int>>> Query;
    vector<int> ans(Q, -1);
    for(int i = 0; i < Q; ++i){
        int A, B; cin >> A >> B, --A, --B;
        int sa = S[A], sb = S[B];
        if(sa == sb){
            ans[i] = uf.Same(A, B);
        }
        else{
            if(sa > sb) swap(sa, sb);
            Query[{sa, sb}].push_back({i, A, B});
        }
    }
    uf.Record();
    for(auto [state_pair, qs] : Query){
        for(auto [u, v] : Road[state_pair]){
            uf.Unite(u, v);
        }
        for(auto [i, a, b] : qs){
            ans[i] = uf.Same(a, b);
        }
        uf.Rollback();
    }
    for(auto a : ans){
        assert(a != -1);
        cout << a << '\n';
    }
}

#line 1 "verify/AOJ-0723.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/challenges/sources/JOI/Prelim/0723"

#line 2 "Library/Common.hpp"

/**
 * @file Common.hpp
 */

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
using namespace std;

using ll = int64_t;
using ull = uint64_t;

constexpr const ll INF = (1LL << 62) - (3LL << 30) - 1;
#line 2 "Library/DataStructure/RollbackUnionFind.hpp"

class RollbackUnionFind{
    public:
    RollbackUnionFind(size_t n) : data_(n, -1), record_(0){}

    int Find(const int k) const {
        if(data_[k] < 0) return k;
        return Find(data_[k]);
    }

    bool Same(const int x, const int y) const {
        return Find(x) == Find(y);
    }

    bool Unite(int x, int y){
        x = Find(x), y = Find(y);
        history_.push({x, data_[x]});
        history_.push({y, data_[y]});
        if(x == y) return false;
        if(data_[x] > data_[y]) swap(x, y);
        data_[x] += data_[y];
        data_[y] = x;
        return true;
    }
    
    int Size(const int k) const {
        return -data_[Find(k)];
    }
    
    vector<vector<int>> Group() const {
        vector<vector<int>> ret(data_.size());
        for(int i = 0; i < data_.size(); ++i){
            ret[Find(i)].emplace_back(i);
        }
        ret.erase(remove_if(begin(ret), end(ret), [&](vector<int> &v){
            return v.empty();
        }), end(ret));
        return ret;
    }

    void Record(){
        record_ = history_.size();
    }

    void Undo(){
        auto [y, dy] = history_.top();
        history_.pop();
        data_[y] = dy;
        auto [x, dx] = history_.top();
        history_.pop();
        data_[x] = dx;
    }

    void Rollback(){
        int state = record_;
        while(state < (int)history_.size()) Undo();
    }
    
    private:
    vector<int> data_;
    stack<pair<int, int>> history_;
    int record_;
};
#line 4 "verify/AOJ-0723.test.cpp"

int main(){
    cin.tie(0)->sync_with_stdio(false);
    int N, M, K; cin >> N >> M >> K;
    vector<int> U(M), V(M);
    for(int i = 0; i < M; ++i) cin >> U[i] >> V[i], --U[i], --V[i];
    vector<int> S(N);
    for(int i = 0; i < N; ++i) cin >> S[i], --S[i];

    RollbackUnionFind uf(N);
    map<pair<int, int>, vector<pair<int, int>>> Road;
    for(int i = 0; i < M; ++i){
        int u = U[i], v = V[i];
        int su = S[u], sv = S[v];
        if(su == sv) uf.Unite(u, v);
        else{
            if(su > sv) swap(su, sv);
            Road[{su, sv}].push_back({u, v});
        }
    }
    int Q; cin >> Q;
    map<pair<int, int>, vector<tuple<int, int, int>>> Query;
    vector<int> ans(Q, -1);
    for(int i = 0; i < Q; ++i){
        int A, B; cin >> A >> B, --A, --B;
        int sa = S[A], sb = S[B];
        if(sa == sb){
            ans[i] = uf.Same(A, B);
        }
        else{
            if(sa > sb) swap(sa, sb);
            Query[{sa, sb}].push_back({i, A, B});
        }
    }
    uf.Record();
    for(auto [state_pair, qs] : Query){
        for(auto [u, v] : Road[state_pair]){
            uf.Unite(u, v);
        }
        for(auto [i, a, b] : qs){
            ans[i] = uf.Same(a, b);
        }
        uf.Rollback();
    }
    for(auto a : ans){
        assert(a != -1);
        cout << a << '\n';
    }
}

Back to top page