一、图的基本实现
import heapq from typing import List, Tuple, Dict, Set class Graph: def __init__(self, vertices: int): """ 初始化图 Args: vertices: 顶点数量 """ self.V = vertices self.graph = [[0] * vertices for _ in range(vertices)] def add_edge(self, u: int, v: int, w: int): """ 添加边 Args: u: 起始顶点 v: 结束顶点 w: 权重 """ self.graph[u][v] = w self.graph[v][u] = w def print_solution(self, parent: List[int]): """ 打印最小生成树 Args: parent: 父节点数组 """ print("边 \t权重") total_weight = 0 for i in range(1, self.V): print(f"{parent[i]} - {i} \t{self.graph[i][parent[i]]}") total_weight += self.graph[i][parent[i]] print(f"最小生成树总权重: {total_weight}")二、普利姆(Prim)算法实现
2.1 基于邻接矩阵实现
class PrimMST: def prim_mst_adj_matrix(self, graph: List[List[int]]) -> Tuple[List[int], int]: """ 普利姆算法（邻接矩阵版本） Args: graph: 邻接矩阵 Returns: 父节点数组和总权重 """ V = len(graph) # 初始化 key = [float('inf')] * V # 存储每个顶点的最小权重 parent = [-1] * V # 存储最小生成树的父节点 mst_set = [False] * V # 记录顶点是否已加入MST # 从第一个顶点开始 key[0] = 0 for _ in range(V): # 选择key值最小的未访问顶点 u = self._min_key(key, mst_set) mst_set[u] = True # 更新相邻顶点的key值 for v in range(V): if (graph[u][v] > 0 and # 存在边 not mst_set[v] and # 未访问 graph[u][v] < key[v]): # 权重更小 key[v] = graph[u][v] parent[v] = u # 计算总权重 total_weight = sum(key[1:]) return parent, total_weight def _min_key(self, key: List[float], mst_set: List[bool]) -> int: """找到最小key值的顶点索引""" min_val = float('inf') min_index = -1 for v in range(len(key)): if not mst_set[v] and key[v] < min_val: min_val = key[v] min_index = v return min_index2.2 邻接表+优先队列优化版本
class PrimMSTOptimized: def prim_mst_adj_list(self, edges: List[List[Tuple[int, int]]]) -> Tuple[List[Tuple[int, int, int]], int]: """ 普利姆算法优化版本（邻接表+优先队列） Args: edges: 邻接表，edges[u] = [(v, weight), ...] Returns: 最小生成树的边列表和总权重 """ V = len(edges) visited = [False] * V min_heap = [] mst_edges = [] total_weight = 0 # 从顶点0开始 heapq.heappush(min_heap, (0, 0, -1)) # (权重, 当前顶点, 父顶点) while min_heap and len(mst_edges) < V: weight, u, parent = heapq.heappop(min_heap) if visited[u]: continue visited[u] = True # 如果不是起始顶点，添加到MST中 if parent != -1: mst_edges.append((parent, u, weight)) total_weight += weight # 处理相邻顶点 for neighbor, w in edges[u]: if not visited[neighbor]: heapq.heappush(min_heap, (w, neighbor, u)) return mst_edges, total_weight三、克鲁斯卡尔(Kruskal)算法实现
class DisjointSet: """并查集实现""" def __init__(self, n: int): self.parent = list(range(n)) self.rank = [0] * n def find(self, x: int) -> int: """查找根节点（路径压缩）""" if self.parent[x] != x: self.parent[x] = self.find(self.parent[x]) return self.parent[x] def union(self, x: int, y: int) -> bool: """合并两个集合（按秩合并）""" root_x = self.find(x) root_y = self.find(y) if root_x == root_y: return False if self.rank[root_x] < self.rank[root_y]: self.parent[root_x] = root_y elif self.rank[root_x] > self.rank[root_y]: self.parent[root_y] = root_x else: self.parent[root_y] = root_x self.rank[root_x] += 1 return True class KruskalMST: def kruskal_mst(self, edges: List[Tuple[int, int, int]], V: int) -> Tuple[List[Tuple[int, int, int]], int]: """ 克鲁斯卡尔算法 Args: edges: 边列表，每个元素为(u, v, weight) V: 顶点数量 Returns: 最小生成树的边列表和总权重 """ # 按权重排序 edges.sort(key=lambda x: x[2]) ds = DisjointSet(V) mst_edges = [] total_weight = 0 for u, v, weight in edges: # 如果不会形成环，则加入MST if ds.union(u, v): mst_edges.append((u, v, weight)) total_weight += weight # 当有V-1条边时，MST已形成 if len(mst_edges) == V - 1: break return mst_edges, total_weight四、测试和比较
def create_test_graph(): """创建测试图""" V = 5 g = Graph(V) # 添加边 edges = [ (0, 1, 2), (0, 3, 6), (1, 2, 3), (1, 3, 8), (1, 4, 5), (2, 4, 7), (3, 4, 9) ] for u, v, w in edges: g.add_edge(u, v, w) return g, edges def test_algorithms(): """测试两种算法""" print("=" * 50) print("最小生成树算法测试") print("=" * 50) # 创建测试图 graph_obj, edges_list = create_test_graph() V = graph_obj.V print("\n1. 普利姆算法（邻接矩阵版本）:") prim = PrimMST() parent, weight = prim.prim_mst_adj_matrix(graph_obj.graph) print("最小生成树的边:") for i in range(1, V): print(f" {parent[i]} - {i} (权重: {graph_obj.graph[i][parent[i]]})") print(f"总权重: {weight}") print("\n2. 普利姆算法（优化版本）:") # 创建邻接表 adj_list = [[] for _ in range(V)] for u, v, w in edges_list: adj_list[u].append((v, w)) adj_list[v].append((u, w)) prim_opt = PrimMSTOptimized() mst_edges_prim, weight_prim = prim_opt.prim_mst_adj_list(adj_list) print("最小生成树的边:") for u, v, w in mst_edges_prim: print(f" {u} - {v} (权重: {w})") print(f"总权重: {weight_prim}") print("\n3. 克鲁斯卡尔算法:") kruskal = KruskalMST() mst_edges_kruskal, weight_kruskal = kruskal.kruskal_mst(edges_list, V) print("最小生成树的边:") for u, v, w in mst_edges_kruskal: print(f" {u} - {v} (权重: {w})") print(f"总权重: {weight_kruskal}") print("\n" + "=" * 50) print("算法比较:") print("1. 普利姆算法适合稠密图，时间复杂度 O(V²)") print("2. 普利姆优化版适合稀疏图，时间复杂度 O(E log V)") print("3. 克鲁斯卡尔算法适合稀疏图，时间复杂度 O(E log E)") print("4. 两种算法都能得到相同的最小生成树总权重") def create_random_graph(n: int, density: float = 0.5): """创建随机图用于测试""" import random edges = [] # 创建邻接矩阵确保连通性 graph = [[0] * n for _ in range(n)] # 首先创建一个生成树确保连通性 for i in range(1, n): parent = random.randint(0, i-1) weight = random.randint(1, 100) graph[i][parent] = weight graph[parent][i] = weight edges.append((min(i, parent), max(i, parent), weight)) # 添加额外的边 max_edges = n * (n-1) // 2 additional_edges = int(density * max_edges) - (n-1) for _ in range(additional_edges): u = random.randint(0, n-1) v = random.randint(0, n-1) if u != v and graph[u][v] == 0: weight = random.randint(1, 100) graph[u][v] = weight graph[v][u] = weight edges.append((min(u, v), max(u, v), weight)) return graph, edges def compare_performance(): """比较算法性能""" import time print("\n" + "=" * 50) print("性能比较测试") print("=" * 50) # 创建不同大小的图 sizes = [10, 50, 100] for n in sizes: print(f"\n测试图大小: {n}个顶点") # 创建图 graph_matrix, edges = create_random_graph(n, 0.3) # 创建邻接表 adj_list = [[] for _ in range(n)] for u, v, w in edges: adj_list[u].append((v, w)) adj_list[v].append((u, w)) # 测试普利姆算法（邻接矩阵） prim = PrimMST() start = time.time() prim.prim_mst_adj_matrix(graph_matrix) prim_time = time.time() - start # 测试普利姆算法（优化版） prim_opt = PrimMSTOptimized() start = time.time() prim_opt.prim_mst_adj_list(adj_list) prim_opt_time = time.time() - start # 测试克鲁斯卡尔算法 kruskal = KruskalMST() start = time.time() kruskal.kruskal_mst(edges, n) kruskal_time = time.time() - start print(f"普利姆(矩阵): {prim_time:.6f}秒") print(f"普利姆(优化): {prim_opt_time:.6f}秒") print(f"克鲁斯卡尔: {kruskal_time:.6f}秒") def main(): """主函数""" # 测试基本功能 test_algorithms() # 比较性能（可选） # compare_performance() print("\n" + "=" * 50) print("示例：手动输入图") print("=" * 50) # 手动输入示例 V = 4 print(f"\n创建一个有{V}个顶点的图:") # 初始化图 g = Graph(V) edges_input = [ (0, 1, 10), (0, 2, 6), (0, 3, 5), (1, 3, 15), (2, 3, 4) ] for u, v, w in edges_input: g.add_edge(u, v, w) print(f"添加边: {u} - {v}, 权重: {w}") print("\n使用普利姆算法计算结果:") prim = PrimMST() parent, weight = prim.prim_mst_adj_matrix(g.graph) g.print_solution(parent) if __name__ == "__main__": main()五、运行结果展示
(mlstat) [haichao@node01 imagenet]$ python demo1.py
==================================================
最小生成树算法测试
==================================================1. 普利姆算法（邻接矩阵版本）:
最小生成树的边:
0 - 1 (权重: 2)
1 - 2 (权重: 3)
0 - 3 (权重: 6)
1 - 4 (权重: 5)
总权重: 162. 普利姆算法（优化版本）:
最小生成树的边:
0 - 1 (权重: 2)
1 - 2 (权重: 3)
1 - 4 (权重: 5)
0 - 3 (权重: 6)
总权重: 163. 克鲁斯卡尔算法:
最小生成树的边:
0 - 1 (权重: 2)
1 - 2 (权重: 3)
1 - 4 (权重: 5)
0 - 3 (权重: 6)
总权重: 16==================================================
算法比较:
1. 普利姆算法适合稠密图，时间复杂度 O(V²)
2. 普利姆优化版适合稀疏图，时间复杂度 O(E log V)
3. 克鲁斯卡尔算法适合稀疏图，时间复杂度 O(E log E)
4. 两种算法都能得到相同的最小生成树总权重==================================================
示例：手动输入图
==================================================创建一个有4个顶点的图:
添加边: 0 - 1, 权重: 10
添加边: 0 - 2, 权重: 6
添加边: 0 - 3, 权重: 5
添加边: 1 - 3, 权重: 15
添加边: 2 - 3, 权重: 4使用普利姆算法计算结果:
边 权重
0 - 1 10
3 - 2 4
0 - 3 5
最小生成树总权重: 19
