leetcode-MST(Prim's Alg)

1135. Connecting Cities With Minimum Cost.py

class Solution:
    def minimumCost(self, N: int, connections: List[List[int]]) -> int:
        """
        Problem formulation:
            We could formulate this problem into minimum spanning tree problem. So, our goal turn to if it's possble to construct minimum spanning tree given an undirected graph. If yes, return weight. Otherwise, return -1.

        Thought process:
            We use Prim's algorithm
                Step 0: build the graph adjacency list.
                Step 1: we can choos any arbitrary node as starting node with zero edge. Put the pair into the priority queue. (Here, we take first node as starting node)
                Step 2: It's a greedy aglrithm. At each step, simply selects the minimum weight edge that adds a new node to the current component(tree) until we make all the nodes in the graph included in the the component. In order to do that, we need help of priority queue and visited array. One is efficiently contain nodes that can be connected to current component in a increasing order of edge's weight and another is to record if the node has be processed or not.
                    Step2-1: We add weight we selected to cost if it has not been processed.
                Step3: If there's only one connected components, return weight. Otherwise, reuturn -1.
        
        Time complexity analysis:
            At most, we have m times operation of insertion and extract elelemt from pq. Each node has been processed once as well as. So, time complexity is O(n + m log m)
        Note:
            1. Prim's algorithm is greedy algorithm.
        
        T:O(n + mlogm)
        S:O(m+n) because we use adjacency list represent graph data strucutre.
        """
        graph = collections.defaultdict(list)
        for a, b, w in connections:
            graph[a-1].append((b-1, w))
            graph[b-1].append((a-1, w))
        
        pq = [(0, 0)] # (zero weight, node 0)
        visited = [False for _ in range(N)]
        cost = 0 # weight
        while len(pq) != 0:
            cur_cost, cur_node = heapq.heappop(pq)
            if visited[cur_node]:
                continue
            cost += cur_cost
            visited[cur_node] = True
            for nei_node, w in graph[cur_node]:
                if not visited[nei_node]:
                    heapq.heappush(pq, (w, nei_node))
        return -1 if sum(visited) != N else cost