From 2d2f1881ca275581e7970dd964106a529c12bb41 Mon Sep 17 00:00:00 2001 From: yenru0 Date: Tue, 18 Feb 2025 18:41:31 +0900 Subject: [PATCH] complete correct 1197.py --- zeta_python/completed/1197.py | 68 +++++++++++++++++------------------ 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/zeta_python/completed/1197.py b/zeta_python/completed/1197.py index 226684a..77614bb 100644 --- a/zeta_python/completed/1197.py +++ b/zeta_python/completed/1197.py @@ -3,48 +3,46 @@ import sys input = sys.stdin.readline -def find(P, a) -> int: - node = a - while node != P[node]: - node = P[node] - return node +class MinimumSpanningTree: + def __init__(self, V: int, E: int, edges: list[tuple[int]]): + self.V = V + self.E = E + self.edges: list[tuple[int]] = edges + self.edges.sort(key=lambda x: x[2]) + self.__parents = [i for i in range(self.V + 1)] -def union(P, a, b) -> None: - ra = find(P, a) - rb = find(P, b) - if ra < rb: - P[rb] = ra - else: - P[ra] = rb + def __root(self, x: int) -> int: + node = x + while node != self.__parents[node]: + node = self.__parents[node] + return node + def __union(self, x: int, y: int): + rx = self.__root(x) + ry = self.__root(y) + if rx > ry: + self.__parents[ry] = rx + else: + self.__parents[rx] = ry -def two_of_them(P, a, b) -> bool: - return find(P, a) == find(P, b) + def solve(self): + mst = [] + for e in self.edges: + a = e[0] + b = e[1] + if self.__root(a) == self.__root(b): + continue + mst.append(e) + self.__union(a, b) - -def get_MST(V, E, Es) -> int: - # sort Edges - Es.sort(key=lambda x: x[-1]) - - P = [i for i in range(V + 1)] - cum_edges = [] - for ei in range(E): - s, e, w = Es[ei] - if not two_of_them(P, s, e): - cum_edges.append(w) - union(P, s, e) - if len(cum_edges) == V - 1: + if len(mst) == self.V - 1: break - - return sum(cum_edges) + return sum(map(lambda x: x[2], mst)) if __name__ == "__main__": V, E = map(int, input().split()) - Es = [] - for _ in range(E): - start, end, weight = map(int, input().split()) - Es.append((start, end, weight)) - - print(get_MST(V, E, Es)) + edges = [tuple(map(int, input().split())) for _ in range(E)] + solver = MinimumSpanningTree(V, E, edges) + print(solver.solve())