Leetcode 1168. Optimize Water Distribution in a Village [Python]

因为多了一个well，也就是总要有一个村子需要自己挖井，所以，按照题目的提示，设置0号村来连接到某一个或者是多个给定的村子，他们的edge weight就是挖井的cost。然后就可以顺利的带入到克鲁斯卡算法里，按照edge weight来排序，UnionFind一下就好。当edge的数量等于节点数量-1的时候（0号虚拟节点加n个给定节点，n+1个总共），说明已经全部链接上了，都有水了。

class UnionFind:
    def __init__(self, n):
        self.n = n
        self.pa = [i for i in range(self.n)]
        
    def find(self, a):
        if self.pa[a] != a:
            self.pa[a] = self.find(self.pa[a])
        return self.pa[a]
    
    def union(self, a, b):
        pofa = self.find(a)
        pofb = self.find(b)
        if pofa == pofb:return False
        self.pa[pofb] = pofa
        return True
        
    
class Solution:
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        uf = UnionFind(n+1)
        
        edge = []
        for u,v,w in pipes:
            edge.append((w,u,v))
        for index, w in enumerate(wells):
            edge.append((w,index+1, 0))
            
        edge.sort(key = lambda x:x[0])
        
        cunt = 0
        res = 0
        for w, u, v in edge:
            if uf.union(u,v):
                cunt += 1
                res += w
            if cunt == n:
                break
        return res

然后是剪枝的做法，使用min heap。首先，建立graph，节点对应其可连接节点和weight的二元组，往heap里丢开始节点0，和一个为0的形式花销。随后，当节点数小于给定节点数n+一个虚拟0号节点（处理自己挖well）时，继续while循环，将heap中的节点和weight弹出，遍历过则continue，没见过此节点，丢到meet里。总weigth+当前edge的weight。代表建立了这个edge。随后，依据当前节点，遍历其的所能连接的目的节点，目的节点未出现在meet中的，丢入heap中，重复上述while循环操作。

class Solution:
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        graph = collections.defaultdict(list)
        for u,v,w in pipes:
            graph[u].append((v,w))
            graph[v].append((u,w))
        for index, cost in enumerate(wells):
            graph[index+1].append((0,cost))
            graph[0].append((index+1,cost))
            
        def prime(start):
            heap = [(0, start)]
            meet = set()
            res = 0
            while len(meet) < n+1:
                w, u = heapq.heappop(heap)
                if u in meet:continue
                meet.add(u)
                res += w
                for v, cost in graph[u]:
                    if v not in meet:
                        heapq.heappush(heap, (cost,v))
            return res
        
        return prime(0)