C++ Library for Competitive Programming
/*
* @title データ構造/素集合データ構造/undo 可能 union-find
*
* verification-helper: IGNORE
* verification-helper: PROBLEM https://atcoder.jp/contests/joi2022yo2/tasks/joi2022_yo2_e
*/
#include <algorithm>
#include <iostream>
#include <map>
#include <utility>
#include <vector>
#include "emthrm/data_structure/union-find/undoable_union-find.hpp"
int main() {
int n, m, k;
std::cin >> n >> m >> k;
std::vector<int> u(m), v(m), s(n);
for (int i = 0; i < m; ++i) {
std::cin >> u[i] >> v[i];
--u[i]; --v[i];
}
for (int i = 0; i < n; ++i) {
std::cin >> s[i];
--s[i];
}
emthrm::UndoableUnionFind union_find(n);
std::map<std::pair<int, int>, std::vector<int>> edges;
for (int i = 0; i < m; ++i) {
if (s[u[i]] == s[v[i]]) {
union_find.unite(u[i], v[i]);
} else {
edges[std::minmax(s[u[i]], s[v[i]])].emplace_back(i);
}
}
union_find.snapshot();
int q;
std::cin >> q;
std::map<std::pair<int, int>, std::vector<int>> queries;
std::vector<int> a(q), b(q);
for (int i = 0; i < q; ++i) {
std::cin >> a[i] >> b[i];
--a[i]; --b[i];
queries[std::minmax(s[a[i]], s[b[i]])].emplace_back(i);
}
std::vector<int> ans(q);
for (const auto& sasb_query : queries) {
if (const auto it = edges.find(sasb_query.first); it != edges.end()) {
for (const int id : it->second) {
union_find.unite(u[id], v[id]);
}
}
for (const int id : sasb_query.second) {
ans[id] = union_find.is_same(a[id], b[id]);
}
union_find.rollback();
}
for (int i = 0; i < q; ++i) {
std::cout << ans[i] << '\n';
}
return 0;
}
#line 1 "test/data_structure/union-find/undoable_union-find.test.cpp"
/*
* @title データ構造/素集合データ構造/undo 可能 union-find
*
* verification-helper: IGNORE
* verification-helper: PROBLEM https://atcoder.jp/contests/joi2022yo2/tasks/joi2022_yo2_e
*/
#include <algorithm>
#include <iostream>
#include <map>
#include <utility>
#include <vector>
#line 1 "include/emthrm/data_structure/union-find/undoable_union-find.hpp"
#line 6 "include/emthrm/data_structure/union-find/undoable_union-find.hpp"
namespace emthrm {
struct UndoableUnionFind {
explicit UndoableUnionFind(const int n) : data(n, -1) {}
int root(const int ver) const {
return data[ver] < 0 ? ver : root(data[ver]);
}
bool unite(int u, int v) {
u = root(u);
history.emplace_back(u, data[u]);
v = root(v);
history.emplace_back(v, data[v]);
if (u == v) return false;
if (data[u] > data[v]) std::swap(u, v);
data[u] += data[v];
data[v] = u;
return true;
}
bool is_same(const int u, const int v) const { return root(u) == root(v); }
int size(const int ver) const { return -data[root(ver)]; }
void undo() {
for (int i = 0; i < 2; ++i) {
data[history.back().first] = history.back().second;
history.pop_back();
}
}
void snapshot() { history.clear(); }
void rollback() {
while (!history.empty()) undo();
}
private:
std::vector<int> data;
std::vector<std::pair<int, int>> history;
};
} // namespace emthrm
#line 15 "test/data_structure/union-find/undoable_union-find.test.cpp"
int main() {
int n, m, k;
std::cin >> n >> m >> k;
std::vector<int> u(m), v(m), s(n);
for (int i = 0; i < m; ++i) {
std::cin >> u[i] >> v[i];
--u[i]; --v[i];
}
for (int i = 0; i < n; ++i) {
std::cin >> s[i];
--s[i];
}
emthrm::UndoableUnionFind union_find(n);
std::map<std::pair<int, int>, std::vector<int>> edges;
for (int i = 0; i < m; ++i) {
if (s[u[i]] == s[v[i]]) {
union_find.unite(u[i], v[i]);
} else {
edges[std::minmax(s[u[i]], s[v[i]])].emplace_back(i);
}
}
union_find.snapshot();
int q;
std::cin >> q;
std::map<std::pair<int, int>, std::vector<int>> queries;
std::vector<int> a(q), b(q);
for (int i = 0; i < q; ++i) {
std::cin >> a[i] >> b[i];
--a[i]; --b[i];
queries[std::minmax(s[a[i]], s[b[i]])].emplace_back(i);
}
std::vector<int> ans(q);
for (const auto& sasb_query : queries) {
if (const auto it = edges.find(sasb_query.first); it != edges.end()) {
for (const int id : it->second) {
union_find.unite(u[id], v[id]);
}
}
for (const int id : sasb_query.second) {
ans[id] = union_find.is_same(a[id], b[id]);
}
union_find.rollback();
}
for (int i = 0; i < q; ++i) {
std::cout << ans[i] << '\n';
}
return 0;
}