使用Prim算法在Python中找到MST的程序

最小生成树是一种为连接无向图的所有顶点而选取的无向图的子集。Prim算法是一种常用的寻找MST的算法。以下是一个使用Prim算法在Python中找到MST的程序示例。

程序示例

import heapq

def minimum_spanning_tree(graph, start_vertex):
    # 用来存储最小生成数的节点和边
    mst = set()
    # 用来存储还未被访问的节点
    unvisited_vertices = set(graph)
    # 用来存储从起始节点开始的距离
    distance = {vertex: float('infinity') for vertex in graph}
    distance[start_vertex] = 0

    # 用heapq来自动排序，pop出距离最短的边，并把对应节点加入访问列表
    priority_queue = [(0, start_vertex)]
    heapq.heapify(priority_queue)

    while unvisited_vertices:
        # 取出距离最短的节点和对应距离
        current_distance, current_vertex = heapq.heappop(priority_queue)

        # 如果节点已经访问过了，则跳过此次循环
        if current_distance > distance[current_vertex]:
            continue

        # 把当前节点加入访问列表
        unvisited_vertices.remove(current_vertex)

        for neighbor, weight in graph[current_vertex].items():
            # 如果邻居已经被访问，则跳过此次循环
            if neighbor not in unvisited_vertices:
                continue

            # 如果邻居到起点的距离更小，则更新距离
            new_distance = weight
            if new_distance < distance[neighbor]:
                distance[neighbor] = new_distance
                heapq.heappush(priority_queue, (new_distance, neighbor))
                mst.add((current_vertex, neighbor))

    return mst

将以下测试代码添加到你的代码中，输入给出的图形并测试结果。

graph = {
    'A': {'B': 2, 'D': 6},
    'B': {'A': 2, 'C': 3, 'D': 8, 'E': 5},
    'C': {'B': 3, 'E': 7},
    'D': {'A': 6, 'B': 8, 'E': 9},
    'E': {'B': 5, 'C': 7, 'D': 9}
}

print(minimum_spanning_tree(graph, 'A'))