Weighted Union Find
(structure/weightedUnionFind.hpp)

• View this file on GitHub
• Last update: 2022-12-04 10:55:39+09:00
• Include: #include "structure/weightedUnionFind.hpp"

Weighted Union Find

• 重み付き Union Find
• weight(y) - weight(x) = w となるようにマージする

参考

• 重み付き Union-Find 木とそれが使える問題のまとめ、および、牛ゲーについて

Verified with

• test/structure/weighted_union_find.aoj_V13_1330.test.cpp

Code

#pragma once
#include <vector>
using namespace std;

template <typename T = int>
struct UnionFind {
  vector<int> par;
  vector<int> rank;
  vector<T> diff_weight;

  UnionFind(int n = 1, T SUM_UNITY = 0) { init(n, SUM_UNITY); }

  void init(int n = 1, T SUM_UNITY = 0) {
    par.resize(n);
    rank.resize(n);
    diff_weight.resize(n);
    for (int i = 0; i < n; i++) {
      par[i] = i;
      rank[i] = 0;
      diff_weight[i] = SUM_UNITY;
    }
  }

  int root(int x) {
    if (par[x] == x) {
      return x;
    } else {
      int r = root(par[x]);
      diff_weight[x] += diff_weight[par[x]];
      return par[x] = r;
    }
  }

  T weight(int x) {
    root(x);
    return diff_weight[x];
  }

  T diff(int x, int y) { return weight(y) - weight(x); }

  bool issame(int x, int y) { return root(x) == root(y); }

  bool merge(int x, int y, T w) {
    w += weight(x);
    w -= weight(y);

    x = root(x);
    y = root(y);
    if (x == y) return false;
    if (rank[x] < rank[y]) {
      swap(x, y);
      w = -w;
    }
    if (rank[x] == rank[y]) rank[x]++;
    par[y] = x;
    diff_weight[y] = w;
    return true;
  }
};

#line 2 "structure/weightedUnionFind.hpp"
#include <vector>
using namespace std;

template <typename T = int>
struct UnionFind {
  vector<int> par;
  vector<int> rank;
  vector<T> diff_weight;

  UnionFind(int n = 1, T SUM_UNITY = 0) { init(n, SUM_UNITY); }

  void init(int n = 1, T SUM_UNITY = 0) {
    par.resize(n);
    rank.resize(n);
    diff_weight.resize(n);
    for (int i = 0; i < n; i++) {
      par[i] = i;
      rank[i] = 0;
      diff_weight[i] = SUM_UNITY;
    }
  }

  int root(int x) {
    if (par[x] == x) {
      return x;
    } else {
      int r = root(par[x]);
      diff_weight[x] += diff_weight[par[x]];
      return par[x] = r;
    }
  }

  T weight(int x) {
    root(x);
    return diff_weight[x];
  }

  T diff(int x, int y) { return weight(y) - weight(x); }

  bool issame(int x, int y) { return root(x) == root(y); }

  bool merge(int x, int y, T w) {
    w += weight(x);
    w -= weight(y);

    x = root(x);
    y = root(y);
    if (x == y) return false;
    if (rank[x] < rank[y]) {
      swap(x, y);
      w = -w;
    }
    if (rank[x] == rank[y]) rank[x]++;
    par[y] = x;
    diff_weight[y] = w;
    return true;
  }
};

Back to top page