モチベーション
筆者は弱小プログラマなわけですが、強プログラマになるために競技プログラミング始めることにしました。
やってみるとわかるんですが、結構覚えること多いんですよね。受験で数学を解いていても思うのですが、思考力はもちろん大事なんだけど試験本番で短時間で解答に至るためには引き出しの多さも大事。
ということで主要なアルゴリズムやデータ構造は１つ１つ実装していきます。
第１回となる今回はUnionFindと言われるデータ構造を実装します。
ちなみに競プロをやるモチベが効率的に実務をこなすことにあるので、競プロのみんながPo**Hubよりも利用しているC++ではなく僕の愛用言語Pythonを使っていきます。（C++のが４倍速いとかどっかの記事でみたけど４倍なんて誤差ですよね）
UnionFind
UnionFindはN個のノードがいくつかのグループに分かれているときに
- 任意の２個のノードa, bが属するグループをつなげて１つのグループにする。
- 任意の２個のノードa, bが同じグループに属するのかを判定する
という機能に特化したデータ構造。
いずれもならし計算量log(N)で実装できます。
アルゴリズム
画像はいずれも素敵なスライドより
実装
データ構造
弱小プログラマはまずデータをどうやってもつべきかを知らないわけですが、UnionFindというのは親ノードが誰なのかを覚えておく木構造です。したがってn個のノードをグループ化したければ各ノードの親を保持するリストがあれば良いですね。
class UnionFind:
def __init__(self, n):
self.parents = list(range(n))
...
ちなみに根ノードには親がいないので親が自分自身のノード=根ノードとして表現します。
最初はみんな別々のグループとして初期化するのでみんな根ノードです。
判定
def is_united(self, x, y):
return self._root(x) == self._root(y)
素敵なスライドから引用させていただいた画像にも載っていたように、根ノードが同じかどうかを判定すれば良いわけですから、_rootメソッドを使って根ノードを引っ張ってきてそれらが等しいかどうかを比べれば良いわけです。
コードを書くときは細かいサブアルゴリズムはすでに出来上がった関数があると思って書いていくことで全体のアルゴリズムを手っ取り早く書いてしまうのが一番わかりやすいと思っています。
３分クッキングで「本日はここに出来上がったじゃがいもの燻製があります」っていうアレがなかったらなかなか料理の全体像が見えてこないじゃないですか。
結合
def union(self, x, y):
self.parents[self._root(y)] = self._root(x)
これも素敵なスライドに書いてあった通りで、xの所属グループとyの所属グループを結合したかったら、どっちかの根ノードを相手の根ノードにくっつけてあげればいいだけです。
くっつけると言ってもくっつけたいノードの親を、くっつける先のノードに指定するだけですね。
根ノードを取得する
これまで使ってきたサブアルゴリズムは_rootだけなので、これを実装したら完成なはずです。
def _root(self, x):
if self.parents[x] == x:
return x
else:
root = self._root(self.parents[x])
return root
根ノードをリターンするためには親が自分自身になるノードが現れるまで再起的に親ノードへと遡っていくだけです。UnionFind木のデータ構造では各ノードの親への参照なので簡単ですね。
でも計算量の都合上Tree構造ができるだけフラットな方が、計算量が軽くなるって偉大なスライドに書いてありました。
そこでコードに若干の工夫を加えます。
def _root(self, x):
if self.parents[x] == x:
return x
else:
root = self._root(self.parents[x])
self.parents[x] = root # 親ノードに直接繋ぎ直す。
return root
これでUnionFind完成！
動かしてみる
uf = UnionFind(5)
uf.union(0, 3)
uf.union(1, 3)
print(uf.is_united(0, 1)) # True
おしまい
一緒にエンジニア活動を頑張ってくれる仲間を探してます。LGTMやフォロー、Twitterのフォローしてくれるととても喜びます。