使用Python实现最小生成树算法
最小生成树算法，即求解一个连通带权无向图的生成树，使得这个生成树的权值最小。常见的实现算法包括：Prim算法和Kruskal算法。
以下为Python实现Prim算法的代码：
import heapq from typing import List, Tuple def prim(graph: List[List[Tuple[int,int]]]) -> List[Tuple[int,int]]: """ :param graph: 二维列表，表示邻接表，其中graph[i][j]表示节点i到节点j的边，边权为graph[i][j][1]。 :return: 返回一个列表，表示最小生成树的所有边，每条边用一个二元组(u, v)表示，其中u和v分别为这条边的两个节点。 """ n = len(graph) mst = [] # 最小生成树中的边 visited = [False] * n # 标记每个节点是否被访问 pq = [] # 小根堆，用于存储与当前生成树相邻的节点 # 将第一个节点加入堆中 visited[0] = True for neighbor in graph[0]: heapq.heappush(pq, (neighbor[1], 0, neighbor[0])) # 按照边权从小到大依次将相邻节点加入生成树 while pq: # 取出堆顶元素，并检查是否已经访问过 wt, u, v = heapq.heappop(pq) if visited[v]: continue visited[v] = True mst.append((u, v)) # 将新节点的相邻节点加入堆中 for neighbor in graph[v]: if not visited[neighbor[0]]: heapq.heappush(pq, (neighbor[1], v, neighbor[0])) return mst
以下为Python实现Kruskal算法的代码：
from typing import List, Tuple def kruskal(n: int, edges: List[Tuple[int,int,int]]) -> List[Tuple[int,int]]: """ :param n: 图中节点数 :param edges: 三元组列表，表示所有边，每条边用一个三元组(u, v, w)表示，其中u和v分别为这条边的两个节点，w为边权。 :return: 返回一个列表，表示最小生成树的所有边，每条边用一个二元组(u, v)表示，其中u和v分别为这条边的两个节点。 """ parent = list(range(n)) # 使用并查集维护节点之间的连通性 edges.sort(key=lambda x: x[2]) # 按边权从小到大排序 mst = [] # 最小生成树中的边 # 遍历所有边，将边加入生成树中 for u, v, w in edges: pu, pv = find(parent, u), find(parent, v) if pu == pv: # 若两个节点已经连通，则跳过 continue parent[pu] = pv # 将两个节点合并为一个连通块 mst.append((u, v)) if len(mst) == n - 1: # 若加入的边数量等于节点数量减一，则生成树已经构建完成 break return mst def find(parent: List[int], i: int) -> int: """并查集查找""" if parent[i] != i: parent[i] = find(parent, parent[i]) return parent[i]
使用方法示例：
graph = [ [(1, 2), (2, 1), (3, 3)], [(0, 2), (2, 4), (4, 5)], [(0, 1), (1, 4), (3, 2), (4, 1)], [(0, 3), (2, 2), (4, 4)], [(1, 5), (2, 1), (3, 4)] ] mst1 = prim(graph) print(mst1) # [(0, 1), (1, 4), (0, 2), (2, 3)] mst2 = kruskal(5, [(i, j, w) for i, neighbors in enumerate(graph) for j, w in neighbors if j > i]) print(mst2) # [(0, 1), (0, 2), (2, 4), (2, 3)]
相关文章