データ構造/素集合データ構造/undo 可能 union-find
(test/data_structure/union-find/undoable_union-find.test.cpp)

• View this file on GitHub
• Last update: 2023-02-26 01:10:01+09:00
• Problem: https://atcoder.jp/contests/joi2022yo2/tasks/joi2022_yo2_e

Depends on

• undo 可能 union-find (include/emthrm/data_structure/union-find/undoable_union-find.hpp)

Code

/*
 * @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;
}

Back to top page