1557. Minimum Number of Vertices to Reach All Nodes

https://leetcode.com/problems/min-cost-to-connect-all-points/

Python

Kruskal's Algorithm

• Time: O(N^2 * logN)
• Space: O(N^2)

class UnionFind:
    def __init__(self, size: int) -> None:
        self.group = [0] * size
        self.rank = [0] * size

        for i in range(size):
            self.group[i] = i

    def find(self, node: int) -> int:
        if self.group[node] != node:
            self.group[node] = self.find(self.group[node])
        return self.group[node]

    def join(self, node1: int, node2: int) -> bool:
        group1 = self.find(node1)
        group2 = self.find(node2)

        # node1 and node2 already belong to same group.
        if group1 == group2:
            return False

        if self.rank[group1] > self.rank[group2]:
            self.group[group2] = group1
        elif self.rank[group1] < self.rank[group2]:
            self.group[group1] = group2
        else:
            self.group[group1] = group2
            self.rank[group2] += 1

        return True

class Solution:
    def minCostConnectPoints(self, points: List[List[int]]) -> int:
        n = len(points)

        edges = []
        for i in range(n):
            for j in range(i+1, n):
                weight = abs(points[i][0]-points[j][0]) + abs(points[i][1]-points[j][1])
                edges.append((weight, i, j))

        edges.sort(key=lambda edge: edge[0])

        uf = UnionFind(n)
        cost = 0
        edges_used = 0

        for weight, i, j in edges:
            if uf.join(i, j):
                cost += weight
                edges_used += 1

                if edges_used == n-1:
                    break
        return cost