たこすの競プロライブラリメモ

C++の個人的競プロライブラリです。99%拝借。

UnionFind木 (Union-Find tree)

目的

要素xと要素yが同じグループに属するかどうかを判定するためのデータ構造。
また、要素xと要素yそれぞれが属するグループを併合することが可能。

関数

root(x):xの根を返す。
unite(x,y):xとyを併合する。
same(x,y):xとyが同じ木に属するならtrue,そうでなければfalseを返す。

コード

struct UnionFind
{
  vector<long long> par; // par[i]:iの親の番号 (例) par[3] = 2 : 3の親が2

  UnionFind(int N) : par(N)
  { //最初は全てが根であるとして初期化
    for (int i = 0; i < N; i++)
      par[i] = i;
  }

  long long root(long long x)
  { // データxが属する木の根を再帰で得る:root(x) = {xの木の根}
    if (par[x] == x)
      return x;
    return par[x] = root(par[x]);
  }

  void unite(long long x, long long y)
  {                   // xとyの木を併合
    long long rx = root(x); //xの根をrx
    long long ry = root(y); //yの根をry
    if (rx == ry)
      return;     //xとyの根が同じ(=同じ木にある)時はそのまま
    par[rx] = ry; //xとyの根が同じでない(=同じ木にない)時:xの根rxをyの根ryにつける
  }

  bool same(long long x, long long y)
  { // 2つのデータx, yが属する木が同じならtrueを返す
    long long rx = root(x);
    long long ry = root(y);
    return rx == ry;
  }
};

コードテスト

#include <bits/stdc++.h>
using namespace std;

struct UnionFind
{
  vector<long long> par; // par[i]:iの親の番号 (例) par[3] = 2 : 3の親が2

  UnionFind(int N) : par(N)
  { //最初は全てが根であるとして初期化
    for (int i = 0; i < N; i++)
      par[i] = i;
  }

  long long root(long long x)
  { // データxが属する木の根を再帰で得る:root(x) = {xの木の根}
    if (par[x] == x)
      return x;
    return par[x] = root(par[x]);
  }

  void unite(long long x, long long y)
  {                   // xとyの木を併合
    long long rx = root(x); //xの根をrx
    long long ry = root(y); //yの根をry
    if (rx == ry)
      return;     //xとyの根が同じ(=同じ木にある)時はそのまま
    par[rx] = ry; //xとyの根が同じでない(=同じ木にない)時:xの根rxをyの根ryにつける
  }

  bool same(long long x, long long y)
  { // 2つのデータx, yが属する木が同じならtrueを返す
    long long rx = root(x);
    long long ry = root(y);
    return rx == ry;
  }
};

int main()
{
  ios::sync_with_stdio(false);
  cin.tie(0);

  UnionFind tree(6);

  tree.unite(1, 2);//1と2が同じグループ
  tree.unite(3, 4);//3と4が同じグループ
  tree.unite(3, 5);//3と5が同じグループ

  if (tree.same(1, 2)==true){ //1と2が同じか判定
    cout << "1-2:Yes" << endl;
  }
  else{
    cout << "1-2:No" << endl;
  }

  if (tree.same(1, 3) == true){ //1と3が同じか判定
    cout << "1-3:Yes" << endl;
  }
  else{
    cout << "1-3:No" << endl;
  }

  if (tree.same(4, 5) == true){ //4と5が同じか判定
    cout << "4-5:Yes" << endl;
  }
  else  {
    cout << "4-5:No" << endl;
  }
}
1-2:Yes
1-3:No
4-5:Yes

出典

qiita.com