ABC350Dを解いた【UnionFind】

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

昨日勉強したUnionFindを早速使いこなすために

解説でUnionFindでも解けると見た記憶のある
この問題を解いていく

実装コード

UnionFind使って各グループごとの頂点数を求める。

あとは解説の通りに答えを出すだけ(以下引用)

頂点数 $n$ の連結成分には最大で $\frac{n(n-1)}{2}$ の辺が存在できるため、求める答えは、この値の連結成分ごとの和からもともと存在する辺の個数を引いたものです。

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, M = rLI()
    
    UF = UnionFind(N)
    
    for _ in range(M):
        a,b = rLI1()
        UF.union(a,b)
    
    ans = 0
    for i in range(N):
        if UF.parent[i] == -1:
            z = UF.rank[i]
            # err(i,z,~-z*z//2)
            ans += (~-z*z)//2
    ans -= M
    print(ans)
if __name__ == '__main__':
    main()

まとめ

UnionFindを使ってABC350のD問題を解いた

解説みてベタ移植で丸写しの解答だが
なにかの力になっている…ハズ

余談

昨日作ったUnionFindにバグが有って30分ハマったorz

位数がイコールの場合に何故か+1しかしていなかったので合計が合わなくなっていた
(昨日の記事も修正済)

修正内容

main.py

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
-       elif rx > ry:
-           self.parent[yr] = xr
-           self.rank[xr] += ry
        else:
            self.parent[yr] = xr
-           self.rank[xr] += 1
+           self.rank[xr] += ry
    def same(self,x,y):
        # err(self.root(x), self.root(y))
        return self.root(x) == self.root(y)