连接所有点的最小费用 -- Kruskal算法应用

1584. 连接所有点的最小费用

class Solution:
    """
    Kruskal 算法应用
    1584. 连接所有点的最小费用
    https://leetcode.cn/problems/min-cost-to-connect-all-points/
    """
    def minCostConnectPoints(self, points: List[List[int]]) -> int:
        n = len(points)
        # 生成两两节点之间的边权重
        edges = []
        for i in range(n):
            for j in range(i, n):
                edges.append([i, j, abs(points[j][0] - points[i][0]) + abs(points[j][1] - points[i][1])])
        
        edges = sorted(edges, key=lambda x: x[2])
        # 执行 kruskal 算法
        mst = 0
        uf = UnionFind2(n)
        for edge in edges:
            u = edge[0]
            v = edge[1]
            cost = edge[2]
            
            if not uf.connected(u, v):
                mst += cost
                uf.union(u, v)
                
        return mst

class UnionFind2:
    """
    并查集算法
    """
    def __init__(self, n):
        self.parent = [i for i in range(n)]
        self.rank = [0] * n
        self.count = n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])

        return self.parent[x]

    def union(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)

        if root_u == root_v:
            return

        if self.rank[u] > self.rank[v]:
            self.parent[root_v] = root_u
        elif self.rank[u] < self.rank[v]:
            self.parent[root_u] = root_v
        else:
            self.parent[root_v] = root_u
            self.rank[root_u] += 1

        self.count -= 1

    def connected(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)

        return root_v == root_u