ds/DSU/DSU.cpp

• View this file on GitHub
• Last update: 2026-01-31 18:57:43+08:00

constructor

DSU(int n): create a disjoint set of size $\text{n}$ without additional data attacted

DSU(vector<T> init, F f(T &larger_set, T &smaller_set) -> void): create a disjoint set where $\text{i}$-th set is attached with $\text{init[i]}$. when calling merge, data would be merged from smaller set to larger set using function $\text{f}$.

example usage:

standard DSU: https://judge.yosupo.jp/submission/349238

DSU attach with $O(1)$ data: https://codeforces.com/edu/course/2/lesson/7/1/practice/contest/289390/submission/360830642

DSU attach with $O(\text{size})$ data and use small-to-large style merge: https://codeforces.com/contest/600/submission/360850419

Verified with

• test/incremental_scc.test.cpp
• test/minimum_spanning_tree_Kruskal.test.cpp
• test/minimum_spanning_tree_Prim.test.cpp
• test/unionfind.test.cpp

Code

template<class T = int, typename F = void*>
struct DSU {
  vi sz_par;
  vc<T> data;
  F op;

  DSU(int n) requires same_as<F, void*> : sz_par(n, -1), op(nullptr) {}

  DSU(vc<T> init, F f) requires invocable<F, T&, T&> &&
    (!invocable<F, T, T&>) && (!invocable<F, T&, T>)
    : sz_par(std::size(init), -1), data(std::move(init)), op(f) {}

  int query(int v) {
    int r = v;
    while(sz_par[r] >= 0) r = sz_par[r];
    while(v != r) {
      int nxt = sz_par[v];
      sz_par[v] = r, v = nxt;
    }
    return r;
  }

  bool merge(int v1, int v2) {
    int b1 = query(v1), b2 = query(v2);

    if (b1 == b2)
      return false;

    if (-sz_par[b1] > -sz_par[b2])
      swap(b1, b2);

    sz_par[b2] += sz_par[b1];
    sz_par[b1] = b2;
    if constexpr (!same_as<F, void*>)
      op(data[b2], data[b1]);

    return true;
  }

  int size(int v) { return v = query(v), -sz_par[v]; }
  const T& get(int v) requires (!same_as<F, void*>) {
    return data[query(v)];
  }
};

#line 1 "ds/DSU/DSU.cpp"
template<class T = int, typename F = void*>
struct DSU {
  vi sz_par;
  vc<T> data;
  F op;

  DSU(int n) requires same_as<F, void*> : sz_par(n, -1), op(nullptr) {}

  DSU(vc<T> init, F f) requires invocable<F, T&, T&> &&
    (!invocable<F, T, T&>) && (!invocable<F, T&, T>)
    : sz_par(std::size(init), -1), data(std::move(init)), op(f) {}

  int query(int v) {
    int r = v;
    while(sz_par[r] >= 0) r = sz_par[r];
    while(v != r) {
      int nxt = sz_par[v];
      sz_par[v] = r, v = nxt;
    }
    return r;
  }

  bool merge(int v1, int v2) {
    int b1 = query(v1), b2 = query(v2);

    if (b1 == b2)
      return false;

    if (-sz_par[b1] > -sz_par[b2])
      swap(b1, b2);

    sz_par[b2] += sz_par[b1];
    sz_par[b1] = b2;
    if constexpr (!same_as<F, void*>)
      op(data[b2], data[b1]);

    return true;
  }

  int size(int v) { return v = query(v), -sz_par[v]; }
  const T& get(int v) requires (!same_as<F, void*>) {
    return data[query(v)];
  }
};

Back to top page