题解 | 最小生成树

#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
# 返回最小的花费代价使得这n户人家连接起来
# @param n int整型 n户人家的村庄
# @param m int整型 m条路
# @param cost int整型二维数组 一维3个参数，表示连接1个村庄到另外1个村庄的花费的代价
# @return int整型
#
class Solution:
    def miniSpanningTree(self , n: int, m: int, cost: List[List[int]]) -> int:
        # write code here
        # kruskal算法+并查集
        parent = list(range(n+1))
        rank = [0]*(n+1)

        def find(x:int) -> int:
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]
        
        def union(x:int,y:int) -> bool:
            root_x = find(x)
            root_y = find(y)
            if root_x == root_y:
                return False
            if rank[root_x] < rank[root_y]:
                parent[root_x] = root_y
            elif rank[root_x] > rank[root_y]:
                parent[root_y] = root_x
            else:
                parent[root_y] = root_x
                rank[root_x] += 1
            return True
        
        cost.sort(key=lambda x:x[2])
        total_weight = 0
        edges_count = 0

        for u,v,w in cost:
            if union(u,v):
                total_weight += w
                edges_count += 1
                if edges_count == n-1:
                    break
        
        return total_weight