1135. Connecting Cities With Minimum Cost

import heapq
class Solution:
    def minimumCost(self, n: int, connections: List[List[int]]) -> int:
        """
        main idea: MST (Prim)
        tc O((E+V)*lgE + E)  sc O(E+V)
        1. build graph based on connections 
        2. create a priority queue, put any start point along with its cost in the queue to start 
        3. pop off one node, check if its not visited, update it to total cost 
        4. check current node's neighbour, if neighbour is not seen, push onto PQ 
        early termination: if total visited nodes == n, break
        """
        g = collections.defaultdict(list)
        seen = set()
        for u,v,w in connections:
            g[u].append((w,v))
            g[v].append((w,u))
        cost = 0
        # cnt = 0
        q = [(0,n)]# (cost, node_id)
        #heapq.heappush(q,(0,n))
        while q:
            w,node = heapq.heappop(q)
            if node not in seen:
                seen.add(node)
                cost += w
                #cnt += 1
                if len(seen) >= n :
                    break
                for wei,nei in g[node]:
                    # pruning, reduce amount of vertices that already added into spanning tree
                    if nei not in seen:   
                        heapq.heappush(q,(wei,nei)) # tc O(VlgE)
        return cost if len(seen) == n else -1

class UF:
    def __init__(self,n):
        self.parent = list(range(n))
        self.rank = [1] *n
        self.size = n

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

    def union(self, u, v):
        pu,pv = self.find(u), self.find(v)
        if pu == pv:
            return False
        if self.rank[pv] > self.rank[pu]:
            self.parent[pu] = pv
            self.rank[pu] += self.rank[pv]
        else:
            self.parent[pv] = pu
            self.rank[pu] += self.rank[pv]
        return True


class Solution:
    def minimumCost(self, n: int, connections: List[List[int]]) -> int:
     """
     tc O(ElgE + E)  sc O(V)
     1. sort the edges by weight
     2. for each edge, find if two vertices belong to one UF, if not, join two points
     3. each time union two vertices, reduce the uf size by one  
     4. early termination: if uf.size == 1, break 
     """
        uf = UF(n)
        connections.sort(key = lambda x: x[2])
        cost = 0
        for u,v,w in connections:
            if uf.union(u-1,v-1):
                cost += w
                uf.size -=1
                if uf.size == 1:
                    break
        return cost if uf.size == 1 else -1