实现最小生成树计算的Prim算法和Kruskal算法python

Prim算法实现：

from heapq import heappush, heappop # 堆优化
from collections import defaultdict

def prim(graph, start):
    mst = defaultdict(set) # 存储最小生成树的结果
    visited = set([start]) # 已经访问过的节点
    edges = [
        (cost, start, to)
        for to, cost in graph[start].items()
    ] # 从起始节点开始的所有边

    # 将边按照权值排序
    edges.sort()

    # 遍历每一条边
    while edges:
        cost, frm, to = heappop(edges) # 取出当前最小的边
        if to not in visited: # 如果该节点未访问过，加入最小生成树
            visited.add(to)
            mst[frm].add(to)
            for to_next, cost in graph[to].items(): # 将新加入的节点的边加入边集中
                if to_next not in visited:
                    heappush(edges, (cost, to, to_next))

    return mst

Kruskal算法实现：

def kruskal(graph):
    mst = defaultdict(set) # 存储最小生成树的结果
    vertices = set(graph.keys()) # 所有节点
    edges = [
        (graph[v1][v2], v1, v2)
        for v1 in vertices
        for v2 in graph[v1]
    ] # 所有边

    # 将边按照权值排序
    edges.sort()

    # 遍历每一条边
    for cost, frm, to in edges:
        if to not in vertices: # 如果该边的终点不在当前最小生成树中，加入该边
            continue
        vertices.remove(to)
        mst[frm].add(to)
        mst[to].add(frm)

    return mst

其中，图的数据结构可以使用字典来表示，例如：

graph = {
    'A': {'B': 2, 'D': 6},
    'B': {'A': 2, 'C': 3, 'D': 8},
    'C': {'B': 3},
    'D': {'A': 6, 'B': 8},
}

其中，每个键（例如'A'）表示一个节点，对应的值（例如{'B': 2, 'D': 6}）表示与该节点相连的边及其权值。例如，图中节点'A'与节点'B'之间的边权为2，节点'A'与节点'D'之间的边权为6。