python实现prim算法

Prim算法是一种求解最小生成树的贪心算法，以下是Python实现Prim算法的示例代码：

import heapq

def prim(graph):
    # 初始化
    visited = set()
    mst = []
    start_vertex = list(graph.keys())[0]
    heap = [(0, start_vertex, start_vertex)]
    # 执行Prim算法
    while heap:
        weight, curr_vertex, prev_vertex = heapq.heappop(heap)
        if curr_vertex not in visited:
            visited.add(curr_vertex)
            mst.append((prev_vertex, curr_vertex, weight))
            for neighbor, weight in graph[curr_vertex].items():
                if neighbor not in visited:
                    heapq.heappush(heap, (weight, neighbor, curr_vertex))
    # 返回最小生成树
    return mst

该算法的输入为一个字典类型的图，其中键为顶点，值为另一个字典，表示该顶点与它相邻的顶点以及对应的边权重。例如，对于下图所示的无向图：

A --4-- B
| \    |
2  5   1
|   \  |
D --3-- C

可以用以下代码表示：

graph = {
    'A': {'B': 4, 'D': 2},
    'B': {'A': 4, 'C': 1, 'D': 5},
    'C': {'B': 1, 'D': 3},
    'D': {'A': 2, 'B': 5, 'C': 3}
}

输出为一个列表，表示最小生成树中的所有边和对应的权重。例如，对于上述图，输出为：

[('A', 'D', 2), ('D', 'C', 3), ('C', 'B', 1)]