test/unittest/union-finds.test.cpp

• View this file on GitHub
• Last update: 2026-08-04 11:59:54+00:00

Depends on

• Persistent Array (structure/others/persistent-array.hpp)
• structure/union-find/bipartite-graph.hpp
• Partially Persistent Union Find (structure/union-find/partially-persistent-union-find.hpp)
• Persistent Union Find (永続 Union Find) (structure/union-find/persistent-union-find.hpp)
• structure/union-find/union-find-undo.hpp
• Union Find (structure/union-find/union-find.hpp)
• structure/union-find/weighted-union-find.hpp

Code

// competitive-verifier: STANDALONE

#include <algorithm>
#include <cassert>
#include <vector>

#include "../../structure/union-find/bipartite-graph.hpp"
#include "../../structure/union-find/partially-persistent-union-find.hpp"
#include "../../structure/union-find/persistent-union-find.hpp"
#include "../../structure/union-find/union-find-undo.hpp"
#include "../../structure/union-find/union-find.hpp"
#include "../../structure/union-find/weighted-union-find.hpp"

int main() {
  {
    UnionFind uf(6);
    assert(uf.unite(0, 1));
    assert(uf.unite(1, 2));
    assert(!uf.unite(0, 2));
    assert(uf.same(0, 2));
    assert(!uf.same(0, 3));
    assert(uf.size(1) == 3);
    uf.unite(3, 4);
    auto groups = uf.groups();
    for (auto& group : groups) std::sort(group.begin(), group.end());
    std::sort(groups.begin(), groups.end());
    assert((groups == std::vector<std::vector<int>>{{0, 1, 2}, {3, 4}, {5}}));
  }

  {
    WeightedUnionFind<long long> uf(4);
    assert(uf.unite(0, 1, 3));
    assert(uf.unite(1, 2, -5));
    assert(uf.diff(0, 2) == -2);
    assert(uf.diff(2, 0) == 2);
    assert(!uf.unite(0, 2, -2));
  }

  {
    UnionFindUndo uf(4);
    uf.unite(0, 1);
    uf.snapshot();
    uf.unite(1, 2);
    assert(uf.size(0) == 3);
    uf.rollback();
    assert(uf.size(0) == 2);
    assert(uf.size(2) == 1);
    uf.unite(2, 3);
    uf.undo();
    assert(uf.size(2) == 1);
  }

  {
    PartiallyPersistentUnionFind uf(5);
    uf.unite(1, 0, 1);
    uf.unite(3, 1, 2);
    uf.unite(5, 3, 4);
    assert(uf.size(0, 0) == 1);
    assert(uf.size(1, 0) == 2);
    assert(uf.size(2, 2) == 1);
    assert(uf.size(3, 0) == 3);
    assert(uf.find(4, 0) == uf.find(4, 2));
    assert(uf.find(4, 0) != uf.find(4, 3));
  }

  {
    PersistentUnionFind original(5);
    auto first = original;
    first.unite(0, 1);
    auto second = first;
    second.unite(1, 2);
    assert(original.size(0) == 1);
    assert(first.size(0) == 2);
    assert(first.size(2) == 1);
    assert(second.size(0) == 3);
  }

  {
    auto add_edge = [](BipartiteGraph& graph, int n, int u, int v) {
      graph.unite(u, v + n);
      graph.unite(u + n, v);
    };

    BipartiteGraph path(3);
    add_edge(path, 3, 0, 1);
    add_edge(path, 3, 1, 2);
    assert(path.bipartite_graph_coloring());
    assert(path[0] != path[1]);
    assert(path[1] != path[2]);

    BipartiteGraph triangle(3);
    add_edge(triangle, 3, 0, 1);
    add_edge(triangle, 3, 1, 2);
    add_edge(triangle, 3, 2, 0);
    assert(!triangle.bipartite_graph_coloring());
  }
}

#line 1 "test/unittest/union-finds.test.cpp"
// competitive-verifier: STANDALONE

#include <algorithm>
#include <cassert>
#include <vector>

#line 2 "structure/union-find/bipartite-graph.hpp"

#line 4 "structure/union-find/bipartite-graph.hpp"

#line 2 "structure/union-find/union-find.hpp"

#line 4 "structure/union-find/union-find.hpp"
#include <cstddef>
#include <utility>
#line 7 "structure/union-find/union-find.hpp"

struct UnionFind {
  std::vector<int> data;

  UnionFind() = default;

  explicit UnionFind(std::size_t sz) : data(sz, -1) {}

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y) return false;
    if (data[x] > data[y]) std::swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return true;
  }

  int find(int k) {
    if (data[k] < 0) return (k);
    return data[k] = find(data[k]);
  }

  int size(int k) { return -data[find(k)]; }

  bool same(int x, int y) { return find(x) == find(y); }

  std::vector<std::vector<int>> groups() {
    int n = (int)data.size();
    std::vector<std::vector<int>> ret(n);
    for (int i = 0; i < n; i++) {
      ret[find(i)].emplace_back(i);
    }
    ret.erase(
        std::remove_if(ret.begin(), ret.end(),
                       [&](const std::vector<int>& v) { return v.empty(); }),
        ret.end());
    return ret;
  }
};
#line 6 "structure/union-find/bipartite-graph.hpp"

struct BipartiteGraph : UnionFind {
  std::vector<int> color;

  BipartiteGraph(int v) : UnionFind(v + v), color(v + v, -1) {}

  bool bipartite_graph_coloring() {
    const int n = (int)color.size() / 2;
    for (int i = 0; i < n; i++) {
      int a = find(i);
      int b = find(i + n);
      if (a == b) return (false);
      if (color[a] < 0) color[a] = 0, color[b] = 1;
    }
    return (true);
  }

  bool operator[](int k) { return (bool(color[find(k)])); }
};
#line 2 "structure/union-find/partially-persistent-union-find.hpp"

#line 4 "structure/union-find/partially-persistent-union-find.hpp"
#include <iterator>
#line 7 "structure/union-find/partially-persistent-union-find.hpp"

struct PartiallyPersistentUnionFind {
  std::vector<int> data;
  std::vector<int> last;
  std::vector<std::vector<std::pair<int, int>>> add;

  PartiallyPersistentUnionFind() {}

  PartiallyPersistentUnionFind(int sz) : data(sz, -1), last(sz, 1e9), add(sz) {
    for (auto& vs : add) vs.emplace_back(-1, -1);
  }

  bool unite(int t, int x, int y) {
    x = find(t, x);
    y = find(t, y);
    if (x == y) return false;
    if (data[x] > data[y]) std::swap(x, y);
    data[x] += data[y];
    add[x].emplace_back(t, data[x]);
    data[y] = x;
    last[y] = t;
    return true;
  }

  int find(int t, int x) {
    if (t < last[x]) return x;
    return find(t, data[x]);
  }

  int size(int t, int x) {
    x = find(t, x);
    return -std::prev(std::lower_bound(add[x].begin(), add[x].end(),
                                       std::make_pair(t, 0)))
                ->second;
  }
};
#line 2 "structure/union-find/persistent-union-find.hpp"

#line 4 "structure/union-find/persistent-union-find.hpp"

#line 2 "structure/others/persistent-array.hpp"

#line 5 "structure/others/persistent-array.hpp"

template <typename T, int LOG>
struct PersistentArray {
  struct Node {
    T data;
    Node* child[1 << LOG] = {};

    Node() {}

    Node(const T& data) : data(data) {}
  };

  Node* root;

  PersistentArray() : root(nullptr) {}

  T get(Node* t, int k) {
    if (k == 0) return t->data;
    return get(t->child[k & ((1 << LOG) - 1)], k >> LOG);
  }

  T get(const int& k) { return get(root, k); }

  std::pair<Node*, T*> mutable_get(Node* t, int k) {
    t = t ? new Node(*t) : new Node();
    if (k == 0) return {t, &t->data};
    auto p = mutable_get(t->child[k & ((1 << LOG) - 1)], k >> LOG);
    t->child[k & ((1 << LOG) - 1)] = p.first;
    return {t, p.second};
  }

  T* mutable_get(const int& k) {
    auto ret = mutable_get(root, k);
    root = ret.first;
    return ret.second;
  }

  Node* build(Node* t, const T& data, int k) {
    if (!t) t = new Node();
    if (k == 0) {
      t->data = data;
      return t;
    }
    auto p = build(t->child[k & ((1 << LOG) - 1)], data, k >> LOG);
    t->child[k & ((1 << LOG) - 1)] = p;
    return t;
  }

  void build(const std::vector<T>& v) {
    root = nullptr;
    for (int i = 0; i < (int)v.size(); i++) {
      root = build(root, v[i], i);
    }
  }
};
#line 6 "structure/union-find/persistent-union-find.hpp"

/*
 * @brief Persistent-Union-Find(永続Union-Find)
 */
struct PersistentUnionFind {
  PersistentArray<int, 3> data;

  PersistentUnionFind() {}

  PersistentUnionFind(int sz) { data.build(std::vector<int>(sz, -1)); }

  int find(int k) {
    int p = data.get(k);
    return p >= 0 ? find(p) : k;
  }

  int size(int k) { return (-data.get(find(k))); }

  bool unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y) return false;
    auto u = data.get(x);
    auto v = data.get(y);

    if (u < v) {
      auto a = data.mutable_get(x);
      *a += v;
      auto b = data.mutable_get(y);
      *b = x;
    } else {
      auto a = data.mutable_get(y);
      *a += u;
      auto b = data.mutable_get(x);
      *b = y;
    }
    return true;
  }
};
#line 2 "structure/union-find/union-find-undo.hpp"

#include <stack>
#line 6 "structure/union-find/union-find-undo.hpp"

struct UnionFindUndo {
  std::vector<int> data;
  std::stack<std::pair<int, int>> history;

  UnionFindUndo(int sz) { data.assign(sz, -1); }

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    history.emplace(x, data[x]);
    history.emplace(y, data[y]);
    if (x == y) return (false);
    if (data[x] > data[y]) std::swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return (true);
  }

  int find(int k) {
    if (data[k] < 0) return (k);
    return (find(data[k]));
  }

  int size(int k) { return (-data[find(k)]); }

  void undo() {
    data[history.top().first] = history.top().second;
    history.pop();
    data[history.top().first] = history.top().second;
    history.pop();
  }

  void snapshot() {
    while (history.size()) history.pop();
  }

  void rollback() {
    while (history.size()) undo();
  }
};
#line 2 "structure/union-find/weighted-union-find.hpp"

#line 5 "structure/union-find/weighted-union-find.hpp"

template <typename T>
struct WeightedUnionFind {
  std::vector<int> data;
  std::vector<T> ws;

  WeightedUnionFind() {}

  WeightedUnionFind(int sz) : data(sz, -1), ws(sz) {}

  int find(int k) {
    if (data[k] < 0) return k;
    auto par = find(data[k]);
    ws[k] += ws[data[k]];
    return data[k] = par;
  }

  T weight(int t) {
    find(t);
    return ws[t];
  }

  bool unite(int x, int y, T w) {
    w += weight(x);
    w -= weight(y);
    x = find(x), y = find(y);
    if (x == y) return false;
    if (data[x] > data[y]) {
      std::swap(x, y);
      w *= -1;
    }
    data[x] += data[y];
    data[y] = x;
    ws[y] = w;
    return true;
  }

  T diff(int x, int y) { return weight(y) - weight(x); }
};
#line 13 "test/unittest/union-finds.test.cpp"

int main() {
  {
    UnionFind uf(6);
    assert(uf.unite(0, 1));
    assert(uf.unite(1, 2));
    assert(!uf.unite(0, 2));
    assert(uf.same(0, 2));
    assert(!uf.same(0, 3));
    assert(uf.size(1) == 3);
    uf.unite(3, 4);
    auto groups = uf.groups();
    for (auto& group : groups) std::sort(group.begin(), group.end());
    std::sort(groups.begin(), groups.end());
    assert((groups == std::vector<std::vector<int>>{{0, 1, 2}, {3, 4}, {5}}));
  }

  {
    WeightedUnionFind<long long> uf(4);
    assert(uf.unite(0, 1, 3));
    assert(uf.unite(1, 2, -5));
    assert(uf.diff(0, 2) == -2);
    assert(uf.diff(2, 0) == 2);
    assert(!uf.unite(0, 2, -2));
  }

  {
    UnionFindUndo uf(4);
    uf.unite(0, 1);
    uf.snapshot();
    uf.unite(1, 2);
    assert(uf.size(0) == 3);
    uf.rollback();
    assert(uf.size(0) == 2);
    assert(uf.size(2) == 1);
    uf.unite(2, 3);
    uf.undo();
    assert(uf.size(2) == 1);
  }

  {
    PartiallyPersistentUnionFind uf(5);
    uf.unite(1, 0, 1);
    uf.unite(3, 1, 2);
    uf.unite(5, 3, 4);
    assert(uf.size(0, 0) == 1);
    assert(uf.size(1, 0) == 2);
    assert(uf.size(2, 2) == 1);
    assert(uf.size(3, 0) == 3);
    assert(uf.find(4, 0) == uf.find(4, 2));
    assert(uf.find(4, 0) != uf.find(4, 3));
  }

  {
    PersistentUnionFind original(5);
    auto first = original;
    first.unite(0, 1);
    auto second = first;
    second.unite(1, 2);
    assert(original.size(0) == 1);
    assert(first.size(0) == 2);
    assert(first.size(2) == 1);
    assert(second.size(0) == 3);
  }

  {
    auto add_edge = [](BipartiteGraph& graph, int n, int u, int v) {
      graph.unite(u, v + n);
      graph.unite(u + n, v);
    };

    BipartiteGraph path(3);
    add_edge(path, 3, 0, 1);
    add_edge(path, 3, 1, 2);
    assert(path.bipartite_graph_coloring());
    assert(path[0] != path[1]);
    assert(path[1] != path[2]);

    BipartiteGraph triangle(3);
    add_edge(triangle, 3, 0, 1);
    add_edge(triangle, 3, 1, 2);
    add_edge(triangle, 3, 2, 0);
    assert(!triangle.bipartite_graph_coloring());
  }
}

Back to top page