584. Min Cost to Connect All Points (Kruskal’s Algorithm)

main.py

class Solution(object):
    def minCostConnectPoints(self, points):
        """
        :type points: List[List[int]]
        :rtype: int
        """
        edges = []
        for i in range(len(points)):
            for j in range(i, len(points)):
                heapq.heappush(edges, (abs(points[i][0] - points[j][0]) + abs(points[i][1] - points[j][1]), i, j))
        
        roots = [i for i in range(len(points))]
        rank = [1] * len(points)
        
        def find(x):
            if roots[x] == x:
                return x
            roots[x] = find(roots[x])
            return roots[x]
        
        def union(n1, n2):
            r1, r2 = find(n1), find(n2)
            if r1 != r2:
                if rank[r1] > rank[r2]:
                    roots[r2] = r1
                elif rank[r2] > rank[r1]:
                    roots[r1] = r2
                else:
                    roots[r1] = r2
                    rank[r2] += 1
                return True
            return False
        c = 0
        num_of_edges = 0
        while num_of_edges < len(points) - 1:
            w, n1, n2 = heapq.heappop(edges)
            if union(n1, n2):
                c += w
                num_of_edges += 1
                
        return c