LeetCode-Python-1168. 水资源分配优化

村里面一共有 n 栋房子。我们希望通过建造水井和铺设管道来为所有房子供水。

对于每个房子 i，我们有两种可选的供水方案：

• 一种是直接在房子内建造水井，成本为 wells[i]；
• 另一种是从另一口井铺设管道引水，数组 pipes 给出了在房子间铺设管道的成本，其中每个 pipes[i] = [house1, house2, cost] 代表用管道将 house1 和 house2 连接在一起的成本。当然，连接是双向的。

请你帮忙计算为所有房子都供水的最低总成本。

提示：

• 1 <= n <= 10000
• wells.length == n
• 0 <= wells[i] <= 10^5
• 1 <= pipes.length <= 10000
• 1 <= pipes[i][0], pipes[i][1] <= n
• 0 <= pipes[i][2] <= 10^5
• pipes[i][0] != pipes[i][1]

思路：

有一点tricky，需要动一动脑筋。

如果供水只能通过铺管道的方法，不能在屋子里建水井，然后问最低成本，那么这道题就是LeetCode-Python-1135. 最低成本联通所有城市，

现在多了一个建造水井的选项，就有了一点变化。

能不能把建造水井也看成是铺设管道呢？

题目给的房子编号是从1 ~ n，那么不妨人为地给地下水标个编号0，这样就可以把水井转化成铺管道：

举例：对于wells = [1, 2, 2]， 就可以看作为pipes[[1, 0, 1], [2, 0, 2], [3, 0, 2]]，

分别代表，从1号房跟地下水连接的开销为1，从2号房跟地下水的连接开销为2，从3号房跟地下水的连接开销为2。

然后题目就变成 LeetCode-Python-1135. 最低成本联通所有城市 了。

class UnionFindSet(object):
    def __init__(self, n):

        self.roots = [i for i in range(n + 1)]
        self.rank = [0 for i in range(n + 1)]
        self.count = n

    def find(self, member):
        tmp = []
        while member != self.roots[member]:
            tmp.append(member)
            member = self.roots[member]
        for root in tmp:
            self.roots[root] = member
        return member

    def union(self, p, q):
        parentP = self.find(p)
        parentQ = self.find(q)
        if parentP != parentQ:
            if self.rank[parentP] > self.rank[parentQ]:
                self.roots[parentQ] = parentP
            elif self.rank[parentP] < self.rank[parentQ]:
                self.roots[parentP] = parentQ
            else:
                self.roots[parentQ] = parentP
                self.rank[parentP] -= 1
            self.count -= 1
            
class Solution(object):
    def minCostToSupplyWater(self, n, wells, pipes):
        """
        :type n: int
        :type wells: List[int]
        :type pipes: List[List[int]]
        :rtype: int
        """
        from heapq import *
        for i, well in enumerate(wells):
            pipes.append([0, i + 1, well])#转井为管
        
        queue = []
        ufs = UnionFindSet(n)
        for start, end, cost in pipes:
            queue.append([cost, start, end]) 
            
        heapify(queue)
        res = 0
        while ufs.count > 0:           
            cost, start, end = heappop(queue)
            if ufs.find(start) == ufs.find(end):
                continue
            res += cost
            ufs.union(end, start)
        return res