競プロ用ライブラリを作る Advent Calendar 2024の3日目です.
UnionFind って?
要素をグループ分けするのに便利なデータ構造です. グラフの連結成分の管理とかに使えます.
- 2つのグループを指定して, そのグループをまとめて1つの大きなグループにする (
union) - 2つの要素を指定して, その2つの要素が同じグループに属しているか判定する (
find)
という操作が$O(\alpha(N))$みたいな計算量でできます. ($\alpha$はアッカーマン関数の逆関数で, 実用上は定数です)
注意点として, 2つのグループを1つにまとめることはできても, 1つのグループを2つに分けたりはできません.
仕組み
各グループを根付き木で表現して, 根の要素を見てグループを判別する, という方針です.
以下, ある要素$x$について, その親を$P(x)$と書くことにします. $x$が根のときは$P(x)$は未定義です.
find
ある要素が属しているグループがどれか, というのを考える時は $x, P(x), P(P(x)), ...$ と根にたどり着くまで探していけばいいです.
union
2つの要素$a$, $b$が属するグループを1つの大きなグループにまとめます. それぞれの根を探して片方の根をもう片方の根の親にしてしまえばいいです.
例えば
こういう木で表現された2つのグループをマージするときは
こういうふうにします. 根の情報しか書き換えなくて良いので愚直な2次元配列での管理にくらべて高速です.
工夫
これだけだと遅い(親まで辿るのに最悪で全部の要素を見る必要がある)ですが, 工夫をすると速くなります.
マージテク
グループをマージするとき, 要素数が大きいほうの根を新しいグループの根にするようにします.
こうじゃなくて
こうするということです
小さい方のグループの要素は全部根から1の距離が1つずつ大きくなります. これはfindにかかるコストが増える, ということで, findのコストが増える要素をできるだけ少なくしたいからどちらのグループが大きいかを考えたわけです.
ただの定数倍高速化と思いきやこれだけで最悪計算量が$O(N)$から$O(\log N)$になります.
経路圧縮
こういう状態があって, $4$の根を探すと, $1$を見つけます.
このとき, 途中に通った$2$, $3$と$4$の根は$1$だと分かったのだから, 直接$1$に繋ぎ直してやると良いです.
こういう風な工夫をすると同じ探索を2回以上しなくなるので, わかりやすく高速化に繋がります.
実装
実装するときは要素ごとに符号付き整数型で持って管理すると良いです.
- 根なら, グループの要素数 × (-1)
- 根じゃないなら, 親の番号
として管理すると別々に値を持つ必要がなくなるので, メモリ量の削減になります. (ac-libraryの実装が参考になります)
それと, 連結成分の数を数えたくなる場面は割と多いのでこの値も管理しておきましょう.
Rustでは値がデフォルトで不変なので, 経路圧縮をしたくてもできないときがあります. 値が不変でもクエリを処理できるように, 経路圧縮しなくても$O(\log N)$でクエリを処理できるので, 妥協して経路圧縮を行わない版のメソッドも作りましょう.
まとめ
いかがでしたか?
明日はこのUnionFindに付加情報を載せる話をやります.