Leetcode 1168. 水资源分配优化

1.题目基本信息

1.1.题目描述

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

对于每个房子 i，我们有两种可选的供水方案：一种是直接在房子内建造水井，成本为 wells[i – 1] （注意 -1 ，因为 索引从0开始 ）；另一种是从另一口井铺设管道引水，数组 pipes 给出了在房子间铺设管道的成本，其中每个 pipes[j] = [house1_j, house2_j, cost_j] 代表用管道将 house1_j 和 house2_j连接在一起的成本。连接是双向的。

请返回 为所有房子都供水的最低总成本 。

1.2.题目地址

https://leetcode.cn/problems/optimize-water-distribution-in-a-village/description/

2.解题方法

2.1.解题思路

最小生成树+Prim算法/Kruskal算法

2.2.解题步骤

第一步，使用邻接表的方式构建图。对于直接在房子内建造水井的情况，创建一个虚拟节点0，将造价作为权值，使其参加到整个无向图中。

第二步，使用Prim算法模板或者Kruskal算法模板解除最小生成树的权值和（详情可以看下代码的注释）

3.解题代码

Python代码（Prim算法）

from collections import defaultdict
import heapq
from typing import Dict,List
# ==> prim算法模板：用于计算无向图的最小生成树及最小权值和
# graph:无向图的邻接表；item项例子：{节点:[[相邻边的权值,相邻边对面的节点],...],...}
# return:最小生成树的权值和；一个合法的最小生成树的边的列表（列表项：[节点,对面节点,两点之间的边的权值]）
def primMinSpanningTree(graph:Dict[object,List[List]]):
    minWeightsSum,treeEdges=0,[]
    firstNode=list(graph.keys())[0]
    # 记录已经加入最小生成树的节点
    visited=set([firstNode])
    # 选择从firstNode开始，相邻的边加到堆中
    neighEdgesHeap=[item+[firstNode] for item in graph[firstNode]]
    heapq.heapify(neighEdgesHeap)
    while len(visited)<len(graph):
        weight,node,node2=heapq.heappop(neighEdgesHeap) # node2为node的weight权值对应的边的对面的节点
        if node not in visited:    # 这个地方必须进行判断，否则会造成重复添加已访问节点，造成最小权值和偏大（因为前面遍历的节点可能将未遍历的共同相邻节点重复添加到堆中）
            minWeightsSum+=weight
            treeEdges.append([node,node2,weight])
            visited.add(node)
            # 遍历新访问的节点的边，加入堆中
            for nextWeight,nextNode in graph[node]:
                if nextNode not in visited:
                    heapq.heappush(neighEdgesHeap,[nextWeight,nextNode,node])
    return minWeightsSum,treeEdges


class Solution:
    # Prim算法
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        # 构建图(邻接表),item项：(费用,房子标记)
        graph=defaultdict(list)
        for index,well in enumerate(wells):
            graph[0].append([well,index+1])
            graph[index+1].append([well,0])
        for pipe in pipes:
            graph[pipe[0]].append([pipe[2],pipe[1]])
            graph[pipe[1]].append([pipe[2],pipe[0]])
        minTotalCost,_=primMinSpanningTree(graph)
        return minTotalCost

Python代码（Kruskal算法）

# # ==> 并查集模板(附优化)
class UnionFind():
    def __init__(self):
        self.roots={}
        # Union优化:存储根节点主导的集合的总节点数
        self.rootSizes={}
    
    def add(self,x):
        if x not in self.roots:
            self.roots[x]=x
            self.rootSizes[x]=1
    
    def find(self,x):
        root=x
        while root != self.roots[root]:
            root=self.roots[root]
        # 优化:压缩路径
        while x!=root:
            temp=self.roots[x]
            self.roots[x]=root
            x=temp
        return root
    
    def union(self,x,y):
        rootx,rooty=self.find(x),self.find(y)
        if rootx!=rooty:
            # 优化:小树合并到大树上
            if self.rootSizes[rootx]<self.rootSizes[rooty]:
                self.roots[rootx]=rooty
                self.rootSizes[rooty]+=self.rootSizes[rootx]
            else:
                self.roots[rooty]=rootx
                self.rootSizes[rootx]+=self.rootSizes[rooty]


class Solution:
    # Kruskal算法
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        # 构建图(邻接表),item项：(费用,房子标记)
        graph=defaultdict(list)
        edgesHeap=[]
        for index,well in enumerate(wells):
            graph[0].append((well,index+1))
            graph[index+1].append((well,0))
            heapq.heappush(edgesHeap,(well,0,index+1))
        for pipe in pipes:
            graph[pipe[0]].append((pipe[2],pipe[1]))
            graph[pipe[1]].append((pipe[2],pipe[0]))
            heapq.heappush(edgesHeap,(pipe[2],pipe[0],pipe[1]))
        # print(graph)
        # 构建点的并查集并初始化节点
        uf=UnionFind()
        for i in range(n+1):
            uf.add(i)
        minTotalCost=0
        # 边的条数固定，一条一条的加
        addedEdgeCnt=0
        while addedEdgeCnt<n:
            # print(edgesHeap)
            cost,house1,house2=heapq.heappop(edgesHeap)
            if uf.find(house1)!=uf.find(house2):
                uf.union(house1,house2)
                minTotalCost+=cost
                addedEdgeCnt+=1
        return minTotalCost

4.执行结果