python最小生成树-prim

Prim算法是一种用于求解图中最小生成树的经典算法，它适用于无向连接的加权图。在Python中实现Prim算法通常会使用优先队列数据结构，如`heapq`模块。以下是Prim算法的基本步骤：

1. **初始选择**：从图中的任意一个顶点开始（通常是权重最低的边），将其添加到已选集合中，并记录其邻接点。

2. **扩展过程**：对于未加入已选集合的每个顶点，检查它与已选集合中的某个顶点相连的一条边，取这条边的权重最小的。将这条边以及与其相连的那个顶点加入到已选集合中。

3. **重复**：重复步骤2，直到所有顶点都被包含在已选集合中，或者无法找到更优的边（这时已选集合构成的就是最小生成树）。

4. **结束**：最后得到的已选集合中的边就是最小生成树的所有边。

import heapq

def prim(graph, start):
    visited = {start: 0}
    edges = []
    for node in graph.keys():
        if node != start:
            edges.append((graph[start][node], node))
    
    heap = [(weight, node) for weight, node in edges]
    heapq.heapify(heap)
    
    min_cost = 0
    while heap:
        weight, node = heapq.heappop(heap)
        if node not in visited:
            visited[node] = weight
            for neighbor, edge_weight in graph[node].items():
                if neighbor not in visited:
                    heapq.heappush(heap, (edge_weight, neighbor))
                    min_cost += edge_weight
    
    return min_cost, visited

# 示例图
graph = {
    'A': {'B': 1, 'C': 5},
    'B': {'A': 1, 'D': 6, 'E': 2},
    'C': {'A': 5, 'D': 1},
    'D': {'B': 6, 'C': 1, 'E': 7, 'F': 9},
    'E': {'B': 2, 'D': 7, 'F': 3},
    'F': {'D': 9, 'E': 3}
}

start_node = 'A'
min_cost, tree = prim(graph, start_node)
print(f"最小生成树总成本: {min_cost}")
print("生成树节点关系:", tree)