UnionFindの勉強【鉄則本A66】

筆者はレート700前後の茶色コーダ

今日は鉄則本のUnionFind問題を解いてみたので
記録としてまとめておく

実装コード

UnionFindはグラフをグループ分けするときに使う木構造

親がマイナス１のときに根とみなす

最初は何も繋がっていないので
大きさ１ですべて親を-1にする

繋げる度にサイズが大きい方を根にする
(そうしたほうが計算量少ないらしい)

根のノードが一緒なら同じグループとみなす

main.py

import sys
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
def err(*args): print(*args, file=sys.stderr)

class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * -~n
        self.rank = [1] * -~n
    def root(self, x):
        while (p:=self.parent[x]) != -1:
            x = p
        return x
    def union(self, x, y):
        xr = self.root(x)
        yr = self.root(y)
        if xr == yr: return
        elif (rx:=self.rank[xr]) < (ry:=self.rank[yr]):
            self.parent[xr] = yr
            self.rank[yr] += rx
        else:
            self.parent[yr] = xr
            self.rank[xr] += ry
    def same(self,x,y):
        # err(self.root(x), self.root(y))
        return self.root(x) == self.root(y)
        

def main():
    N, Q = rLI()
    
    UF = UnionFind(N)
    
    for _ in range(Q):
        q, u, v = rLI()
        # err(q,u,v)
        if q == 1:
            UF.union(u,v)
        if q == 2:
            print("YNeos"[1-UF.same(u,v)::2])
    
    
if __name__ == '__main__':
    main()

まとめ

鉄則本のUnionFind問題を解いた
実装は意外と単純だった

今度はこれを使うABCの過去問とか解いてみたい

おまけ

今回のUnionFind実装をVSCodeスニペットにしたので
使いたい人どうぞ

折りたたみ

python.json

{
    "Union-Find": {
        "prefix": "unionfind",
        "body": [
            "class UnionFind:",
            "    def __init__(self, n):",
            "        self.parent = [-1] * -~n",
            "        self.rank = [1] * -~n",
            "    def root(self, x):",
            "        while (p:=self.parent[x]) != -1:",
            "            x = p",
            "        return x",
            "    def union(self, x, y):",
            "        xr = self.root(x)",
            "        yr = self.root(y)",
            "        if xr == yr: return",
            "        elif (rx:=self.rank[xr]) < (ry:=self.rank[yr]):",
            "            self.parent[xr] = yr",
            "            self.rank[yr] += rx",
            "        else:",
            "            self.parent[yr] = xr",
            "            self.rank[xr] += ry",
            "    def same(self,x,y):",
            "        # err(self.root(x), self.root(y))",
            "        return self.root(x) == self.root(y)"
        ],
        "description": "Union-Findアルゴリズムを実装"
    }
}