クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「2-5 最小全域木問題２（クラスカル法）」(p.101)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 最小全域木問題２（クラスカル法）」(p.101)
のコードをPythonで書き直したものとなっています。

入力


V:頂点数\\E:辺数\\頂点s_iと頂点t_iの間に重みc_iの無向辺が存在する(1\leq{i}\leq{E})


V\,E\\s_1\,t_1\,c_1\\:\\s_E\,t_E\,c_E

入力例

f:id:kuruton456:20201201145758p:plain


7 9
0 1 2
0 4 10
1 2 1
1 3 3
1 5 7
3 5 1
3 6 5
4 5 5
5 6 8


コード

# Union-Find木のコード

class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # xとyが同じ集合に属するか否か
    def same(self, x, y):
        return self.find(x) == self.find(y)


# ここから蟻本のコード

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
es = []

for _ in range(E):
    s, t, c = map(int,input().split())
    es.append([c, s, t])    # c(辺のコスト)でソートできるように最初の要素にしておく

def kruskal():
    es.sort()    # 辺のコストが小さい順にソートする
    UFT = UnionFindTree(V)    # Union-Findの初期化
    res = 0
    for i in range(E):
        e = es[i]
        if not UFT.same(e[1], e[2]):
            UFT.unite(e[1], e[2])
            res += e[0]
    return res

res = kruskal()

print(res)

実行結果

17