complete correct 1197.py

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]
+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 __root(self, x: int) -> int:
+        node = x
+        while node != self.__parents[node]:
+            node = self.__parents[node]
         return node
 
-
-def union(P, a, b) -> None:
-    ra = find(P, a)
-    rb = find(P, b)
-    if ra < rb:
-        P[rb] = ra
+    def __union(self, x: int, y: int):
+        rx = self.__root(x)
+        ry = self.__root(y)
+        if rx > ry:
+            self.__parents[ry] = rx
         else:
-        P[ra] = rb
+            self.__parents[rx] = ry
 
+    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 two_of_them(P, a, b) -> bool:
-    return find(P, a) == find(P, 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())