ds/unionfind/rollback_unionfind.hpp

• View this file on GitHub
• Last update: 2024-12-17 23:15:20+09:00
• Include: #include "ds/unionfind/rollback_unionfind.hpp"

Depends on

• ds/rollback_array.hpp

Verified with

• test/2_library_checker/data_structure/add_remove_query.test.cpp
• test/2_library_checker/graph/dynamic_graph_vertex_add_component_sum.test.cpp

Code

#include "ds/rollback_array.hpp"


struct Rollback_UnionFind {
  int n;
  Rollback_Array<int> dat; // parent or size


  Rollback_UnionFind(int n) : n(n), dat(vc<int>(n, -1)) {}

  int operator[](int v) {
    while (dat.get(v) >= 0) v = dat.get(v);
    return v;
  }

  ll size(int v) { return -dat.get((*this)[v]); }
  int time() { return dat.time(); }
  void rollback(int t) { dat.rollback(t); }
  void reset() { rollback(0); }

  bool merge(int a, int b) {
    a = (*this)[a], b = (*this)[b];
    if (a == b) return false;
    if (dat.get(a) > dat.get(b)) swap(a, b);
    dat.set(a, dat.get(a) + dat.get(b));
    dat.set(b, a);
    return true;
  }
  vc<int> get_all() {
    vc<int> ANS(n);
    FOR(i, n) ANS[i] = (*this)[i];
    return ANS;
  }
};

#line 2 "ds/rollback_array.hpp"

template <typename T>
struct Rollback_Array {
  int N;
  vc<T> dat;
  vc<pair<int, T>> history;

  Rollback_Array() {}
  Rollback_Array(vc<T> x) : N(len(x)), dat(x) {}
  Rollback_Array(int N) : N(N), dat(N) {}
  template <typename F>
  Rollback_Array(int N, F f) : N(N) {
    dat.reserve(N);
    FOR(i, N) dat.eb(f(i));
  }

  int time() { return len(history); }
  void rollback(int t) {
    FOR_R(i, t, time()) {
      auto& [idx, v] = history[i];
      dat[idx] = v;
    }
    history.resize(t);
  }
  T get(int idx) { return dat[idx]; }
  void set(int idx, T x) {
    history.eb(idx, dat[idx]);
    dat[idx] = x;
  }

  vc<T> get_all() {
    vc<T> res(N);
    FOR(i, N) res[i] = get(i);
    return res;
  }
};
#line 2 "ds/unionfind/rollback_unionfind.hpp"

struct Rollback_UnionFind {
  int n;
  Rollback_Array<int> dat; // parent or size


  Rollback_UnionFind(int n) : n(n), dat(vc<int>(n, -1)) {}

  int operator[](int v) {
    while (dat.get(v) >= 0) v = dat.get(v);
    return v;
  }

  ll size(int v) { return -dat.get((*this)[v]); }
  int time() { return dat.time(); }
  void rollback(int t) { dat.rollback(t); }
  void reset() { rollback(0); }

  bool merge(int a, int b) {
    a = (*this)[a], b = (*this)[b];
    if (a == b) return false;
    if (dat.get(a) > dat.get(b)) swap(a, b);
    dat.set(a, dat.get(a) + dat.get(b));
    dat.set(b, a);
    return true;
  }
  vc<int> get_all() {
    vc<int> ANS(n);
    FOR(i, n) ANS[i] = (*this)[i];
    return ANS;
  }
};

Back to top page