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

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

Pythonで理解する蟻本「2-5 Conscription（POJ No.3723）」(p.103)

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

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 Conscription（POJ No.3723）」(p.103)
のコードをPythonで書き直したものとなっています。

入力

N\,M\,R\\x_1\,y_1\,d_1\\:\\x_R\,y_R\,d_R

入力例



5 5 8
4 3 6831
1 3 4583
0 0 6592
0 1 3063
3 3 4975
1 3 2049
4 2 2104
2 2 781


解答

# クラスカル法のコード

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)

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


# ここから蟻本のコード

MAX_R = 50000

# 入力
N, M, R = map(int,input().split())
x = [0] * MAX_R
y = [0] * MAX_R
d = [0] * MAX_R

for i in range(R):
    x[i], y[i], d[i] = map(int,input().split())

V = N + M
E = R
es = []
for i in range(R):
    es.append([-d[i], N + x[i], y[i]])
print(10000 * (N + M) + kruskal())