代码随想录算法训练营第五十七天 | prim算法 kruskal算法

prim算法：

文章链接
题目链接：53.寻宝

思路：

首先要明确prim算法针对的是带权无向图，求得到其最小生成树。最小生成树是，在一个无向带权图中，找到一棵树，使得树的所有边权值之和最小，且包含图中的所有顶点。v个节点，最小生成树中v - 1条边
而本题就是要求带权无向图的最小生成树的总路径长度，采用prim算法。
prim算法的总体思路是：贪心不断将距离最小生成树最近的节点加入生成树中。

• 使用miniDist数组记录节点到最小生成树的最短距离，结束prim算法后，数组中的值就是最小生成树中的权值（一共v - 1条边，在数组中对应下标范围为[2, v]）
• prim三部曲
  • 首先选择不在生成树中，距离最小生成树距离最近的节点
  • 将其加入生成树中
  • 更新非生成树节点到生成树的距离（即miniDist数组）
    同时因为miniDist记录不了节点是否在生成树中，因此使用inTree数组判断节点是否在生成树中
    代码

def main():
    v, e = map(int, input().split())
    graph = [[10001] * (v + 1) for _ in range(v + 1)]   # 节点编号从1开始
    for _ in range(e):
        # prim算法是在无向图中求最小生成树
        v1, v2, val = map(int, input().split())
        graph[v1][v2] = val
        graph[v2][v1] = val
    
    # prim算法
    miniDist = [10001] * (v + 1)    # val最大10000，下标与节点编号相对应
    inTree = [False] * (v + 1)  # 节点是否在最小生成树中
    for _ in range(v):  # 循环v次，将v个节点加入最小生成树
        cur, mindist = -1, float('inf')
        # 找到距离最小生成树最近的节点
        for i in range(1, v + 1):
            if not inTree[i] and miniDist[i] < mindist:
                cur, mindist = i, miniDist[i]   
        # 将其加入最小生成树
        inTree[cur] = True
        # 更新miniDist数组
        for i in range(1, v + 1):
            if not inTree[i] and graph[cur][i] < miniDist[i]:
                miniDist[i] = graph[cur][i]
                
    # 最后miniDist中的值为最小生成树中节点的权值
    # 统计miniDist即得到最小路径总距离
    result = 0
    for i in range(2, v + 1):   # v - 1条边
        result += miniDist[i]
    print(result)
    
if __name__ == '__main__':
    main()

如果想要输出最小生成树的边呢
那么可以参照一下并查集，使用parent数组记录
使用如下：

    miniDist = [10001] * (v + 1)    # val最大10000，下标与节点编号相对应
    inTree = [False] * (v + 1)  # 节点是否在最小生成树中
    parent = [x for x in range(v + 1)]

        # 更新miniDist数组
        for i in range(1, v + 1):
            if not inTree[i] and graph[cur][i] < miniDist[i]:
                miniDist[i] = graph[cur][i]
                parent[i] = cur	# 更新的是parent[i]

---

kruskal算法：

文章链接
题目链接：53.寻宝

思路

kruskal算法也是用于带权无向图得到最小生成树的。
prim算法是维护节点的集合，kruskal算法是维护边的集合
kruskal算法的大致过程为：

• 首先将边按权值顺序排序
• 遍历排序后的边
  • 如果边首尾两个节点在同一个集合中，跳过
  • 如果边首尾两个节点不在同一个集合中，将其加入同一个集合
    需要判断两个节点是否在同一个集合中，那么就是要用到并查集来判断了
    边首尾两个节点不在同一个集合中如下三种情况，以(1, 3)边为例（下面的粗线表示该边已经加入最小生成树中）
    情况1
    情况2
    
    情况3
    
    代码

class Edge():   #  这里是定义了类似C中的结构体Edge
    def __init__(self, l, r, val):
        self.l = l
        self.r = r
        self.val = val
        
class UnionFind():
    def __init__(self, size):
        self.parent = [x for x in range(size + 1)]  # 节点编号从1开始
    def find(self, u):
        if self.parent[u] != u:
            self.parent[u] = self.find(self.parent[u])
        return self.parent[u]
    def join(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)
        if root_u != root_v:
            self.parent[root_v] = root_u
    def is_same(self, u, v):
        return self.find(u) == self.find(v)
        
def kruskal(v, edges):
    edges.sort(key=lambda edge:edge.val)    # 按权值顺序排
    uf = UnionFind(v)
    result_val = 0
    for edge in edges:
        if not uf.is_same(edge.l, edge.r):
            result_val += edge.val
            uf.join(edge.l, edge.r)
    return result_val
    
def main():
    v, e = map(int, input().split())
    edges = []  # 保存边和权值
    for _ in range(e):
        v1, v2, val = map(int, input().split())
        edges.append(Edge(v1, v2, val))
        
    result_val = kruskal(v, edges)
    print(result_val)
if __name__ == '__main__':
    main()

如果想要输出最小生成树的边
因为kruskal算法是从边的角度来求最小生成树的，所以输出边会比prim算法简单

def kruskal(v, edges):
    edges.sort(key=lambda edge:edge.val)    # 按权值顺序排
    result = [] # 存储最小生成树的边
    uf = UnionFind(v)
    result_val = 0
    for edge in edges:
        if not uf.is_same(edge.l, edge.r):
            result.append(edge) # 存储最小生成树的边
            result_val += edge.val
            uf.join(edge.l, edge.r)
            
    # 最后打印显示
    for edge in result:
        print(edge.l, edge.r, edge.val)
    return result_val

---

学习收获：

prim算法和kruskal算法是为了解决带权无向图求最小生成树的问题。
prim算法是画节点的集合，贪心不断将距离生成树最近非生成树的节点加入生成树中，并更新非生成树的节点的距离
kruskal算法是画边的集合，首先对边按权值排序，如果边的首尾两个节点不在同一集合中，就算入最小生成树。其中，不在同一集合中的情况有上面三种，判断两个节点是否在同一集合中用并查集
prim算法按节点来，kruskal算法按边来，因此对于稀疏图，用kruskal算法更好(边数更少)；对于稠密图，用prim算法更好。
prim算法时间复杂度为O(n^2)，其中n为节点数量
kruskal算法时间复杂度为O(nlogn)，其中n为边数