Grid でも思考停止で Union Find がしたい！

はじめに

Grid、つまりは二次元のグラフについての連結判定をしたいときってありますよね。このときに DFS / BFS をサラッとかければいいんですが、実装力がないため、グダグダしてしまいます。

こういうときに役立つ（思考停止で snippet を貼れる）pair<int, int> を素集合として扱える Union Find (Disjoint Set Union) を書いてみました（C++ で類似の記事^[1]が見つからなかったのでこの記事を書いています）

実装

経路圧縮 (Path halving) と Union by size で実装^[2]しています。

普通の vector を用いた Union Find ならば、vector の値へのアクセスの計算量は $O(1)$ であるため、計算量が $O(\alpha(N))$ になります。($\alpha(N)$ はアッカーマン関数の逆関数)

しかし、map の値のアクセスの計算量は $O(\log N)$ であるから、この pair を持たせられる Union Find の計算量は $O(\alpha (N) \log N)$ です。（ここらへんは怪しいので、指摘のほどよろしくお願いします）

初期化にかかる計算量は $O(HW)$ であるため $HW \leq 10^6$ 程度でないと危なそうです。

機能	説明
.root(x)	x を含む集合の代表を返す
.merge(x, y)	x を含む集合と y を含む集合の併合を行う
.same(x, y)	x と y が同じ集合に属するか判定する
.size(x)	x を含む集合の要素数を返す

Code

https://github.com/Kyure-A/competitive-snippets/blob/master/kyre/PairedUnionFind.hpp

例題

ABC 300 C - Cross

Code

class PairedUnionFind
{
public:
  map<pair<int, int>, pair<int, int>> parent;
  map<pair<int, int>, int> set_size;
 
  // constructor
  PairedUnionFind (int h, int w): parent(), set_size()
  {
    for (int i = 0; i < h; ++i)
      {
	for (int j = 0; j < w; ++j)
	  {
	    parent[{i, j}] = {i, j};
	    set_size[{i, j}] = 1;
	  }
      }
  }
 
  pair<int, int> root (pair<int, int> x) // find (path halving)
  {
    while (parent[x] != x)
      {
	parent[x] = parent[parent[x]];
	x = parent[x];
      }
 
    return x;
  }
 
  bool merge (pair<int, int> x, pair<int, int> y) // union by size
  {
    pair<int, int> rx = root(x);
    pair<int, int> ry = root(y);
 
    if (rx == ry) return false;
    else if (set_size[rx] < set_size[ry]) swap(rx, ry); // root(y) のほうがデカいときも merge できるように逆にする
    
    // Operations
    parent[ry] = rx;
    set_size[rx] += set_size[ry];
    return true;
  }
 
  bool same (pair<int, int> x, pair<int, int> y)
  {
    return root(x) == root(y);
  }
 
  int size(pair<int, int> x)
  {
    return set_size[root(x)];
  }
 
};
 
signed main ()
{
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
 
  int h, w; cin >> h >> w;
  vector<vector<char>> c(h, vector<char>(w)); rep(i, h) rep(j, w) cin >> c[i][j];
  PairedUnionFind tree(h, w);
  
  // init
  for (int i = 0; i < h; ++i)
    {
      for (int j = 0; j < w; ++j)
	{
	  if (c[i][j] == '#' and i + 1 < h and j + 1 < w)
	    {
	      if (c[i + 1][j + 1] == '#') tree.merge({i, j}, {i + 1, j + 1});
	    }
 
	  if (c[i][j] == '#' and i + 1 < h and j - 1 >= 0)
	    {
	      if (c[i + 1][j - 1] == '#') tree.merge({i, j}, {i + 1, j - 1});
	    }
 
	      if (c[i][j] == '#' and i - 1 >= 0 and j + 1 < h)
	    {
	      if (c[i - 1][j + 1] == '#') tree.merge({i, j}, {i - 1, j + 1});
	    }
 
	      if (c[i][j] == '#' and i - 1 >= 0 and j - 1 >= 0)
	    {
	      if (c[i - 1][j - 1] == '#') tree.merge({i, j}, {i - 1, j - 1});
	    }
	}
    }
 
  // solve
  set<pair<int, int>> cross;
  
  for (int i = 0; i < h; ++i)
    {
      for (int j = 0; j < w; ++j)
	{
	  bool is_same = false;
	  
	  for (auto x : cross)
	    {
	      if (tree.same({i, j}, x))
		{
		  is_same = true;
		  break;
		}
	    }
 
	  if (is_same or tree.size({i, j}) == 1) continue; // "." の集合を除く
	  
	  cross.insert({i, j});
	}
    }
  
  int n = min(h, w);
  vector<int> answer(n + 1, 0);
 
  for (auto x : cross)
    {
      int size = tree.size(x) / 4; // floor
      answer[size]++;
    }
 
  for (int i = 1; i <= n; ++i)
    {
      cout << answer[i] << " ";
    }
 
  cout << endl;
 
  return 0;
}

競プロ典型 90 問 012 - Red Painting (★4)

Code

class PairedUnionFind
{
public:
  map<pair<int, int>, pair<int, int>> parent;
  map<pair<int, int>, int> set_size;
 
  // constructor
  PairedUnionFind (int h, int w): parent(), set_size()
  {
    for (int i = 0; i < h; ++i)
      {
	for (int j = 0; j < w; ++j)
	  {
	    parent[{i, j}] = {i, j};
	    set_size[{i, j}] = 1;
	  }
      }
  }
 
  pair<int, int> root (pair<int, int> x) // find (path halving)
  {
    while (parent[x] != x)
      {
	parent[x] = parent[parent[x]];
	x = parent[x];
      }
 
    return x;
  }
 
  bool merge (pair<int, int> x, pair<int, int> y) // union by size
  {
    pair<int, int> rx = root(x);
    pair<int, int> ry = root(y);
 
    if (rx == ry) return false;
    
    // Operations
    else if (set_size[rx] < set_size[ry])
      {
	parent[rx] = ry;
	set_size[ry] += set_size[rx];
      }
    else
      {
	parent[ry] = rx;
	set_size[rx] += set_size[ry];
      }
    
    return true;
  }
 
  bool same (pair<int, int> x, pair<int, int> y)
  {
    return root(x) == root(y);
  }
 
  int size(pair<int, int> x)
  {
    return set_size[root(x)];
  }
 
};
 
signed main ()
{
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
 
  int h, w; cin >> h >> w;
  int q; cin >> q;
  
  PairedUnionFind tree(h + 1, w + 1); // 1 indexed
 
  map<pair<int, int>, bool> is_painted;
 
  constexpr int dx[4] = {1, 0, -1, 0};
  constexpr int dy[4] = {0, 1, 0, -1};
  
  while (q--)
    {
      int query; cin >> query;
 
      if (query == 1)
	{
	  int r, c; cin >> r >> c;
	  is_painted[{r, c}] = true;
 
	  for (int i = 0; i < 4; ++i)
	    {
	      int nx = r + dx[i];
	      int ny = c + dy[i];
		  
	      if (1 <= nx and nx <= h and 1 <= ny and ny <= w) // 1 indexed
		{
		  if (is_painted[{nx, ny}]) tree.merge({r, c}, {nx, ny});
		}
	    }
	}
 
      if (query == 2)
	{
	  int r_a, c_a, r_b, c_b; cin >> r_a >> c_a >> r_b >> c_b;
 
	  if (r_a == r_b and c_a == c_b and is_painted[{r_a, c_a}])
	    {
	      cout << "Yes" << endl;
	    }
	  
	  else if ((r_a != r_b or c_a != c_b) and tree.same({r_a, c_a}, {r_b, c_b}))
	    {
	      cout << "Yes" << endl;
	    }
	  
	  else cout << "No" << endl;
	}
    }
  
  return 0;
}

脚注

1. Python で map に ID を振ることで実装した例がありました。冷静に考えるとこの記事にある Union Find の pair<int, int> のみを扱える下位互換ですね。
   https://qiita.com/tomato1997/items/7c001c2a9a1e7f428241

2. これは既存の Union Find を pair<int, int> に書き換えたものなのですが、通常の Union Find を実装する際に参考にした記事です。擬似コード（ほぼ Python）でわかりやすく説明されているのでこれだけ見れば Union Find master! とはいかずとも bachelor くらいにはなれていると思います。
   https://en.wikipedia.org/wiki/Disjoint-set_data_structure