一、图的基本实现
python
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 基于邻接矩阵实现
python
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 邻接表+优先队列优化版本
python
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)算法实现
python
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四、测试和比较
python
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
==================================================
最小生成树算法测试
==================================================
- 普利姆算法（邻接矩阵版本）:
最小生成树的边:
0 - 1 (权重: 2)
1 - 2 (权重: 3)
0 - 3 (权重: 6)
1 - 4 (权重: 5)
总权重: 16
- 普利姆算法（优化版本）:
最小生成树的边:
0 - 1 (权重: 2)
1 - 2 (权重: 3)
1 - 4 (权重: 5)
0 - 3 (权重: 6)
总权重: 16
- 克鲁斯卡尔算法:
最小生成树的边:
0 - 1 (权重: 2)
1 - 2 (权重: 3)
1 - 4 (权重: 5)
0 - 3 (权重: 6)
总权重: 16
==================================================
算法比较:
普利姆算法适合稠密图，时间复杂度 O(V²)
普利姆优化版适合稀疏图，时间复杂度 O(E log V)
克鲁斯卡尔算法适合稀疏图，时间复杂度 O(E log E)
两种算法都能得到相同的最小生成树总权重
==================================================
示例：手动输入图
==================================================
创建一个有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