[LeetCode] 1168. Optimize Water Distribution in a Village

 Problem :  https://leetcode.com/problems/optimize-water-distribution-in-a-village/

This problem can be reduced to a Minimum Spanning Tree problem.

We may consider the cost to build well in a village is the cost to reach to this village from any other villages. So we can use greedy algorithm to get the minimum cost to visit all villages.

class Solution:
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        graph = defaultdict(list)
        
        # consider the cost of laying pipe between village 'a' and 'b'
        # as the cast to reach village 'a' from 'b' and vice versa
        for a, b, cost in pipes:
            graph[a].append((b, cost))
            graph[b].append((a, cost))
        
        visited = set()
        
        heap = []
        total = 0
        
        # consider the cost of building well in one village
        # as the cost to reach to this building from any other building
        for i in range(n):
            heapq.heappush(heap, (wells[i], i+1))
        
        # greedily find the the minimum cost to visit all villages.
        while heap and len(visited) < n:
            cost, a = heapq.heappop(heap)
            
            if a not in visited:
                total += cost
                
                visited.add(a)
                
                for b, bcost in graph[a]:
                    heapq.heappush(heap, (bcost, b))
        
        return total