75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
|
import heapq
|
||
|
|
|
||
|
|
# 图的表示方式
|
||
|
|
class Graph:
|
||
|
|
def __init__(self, vertices):
|
||
|
|
self.V = vertices # 顶点数
|
||
|
|
self.graph = [] # 存储边的列表
|
||
|
|
|
||
|
|
def add_edge(self, u, v, weight):
|
||
|
|
self.graph.append([u, v, weight])
|
||
|
|
|
||
|
|
# Prim 算法
|
||
|
|
def prim_mst(self):
|
||
|
|
mst_set = [False] * self.V # 记录是否已加入 MST
|
||
|
|
edge_heap = [(0, 0, -1)] # (权值, 顶点, 父节点) (-1 表示起始节点无父节点)
|
||
|
|
total_cost = 0
|
||
|
|
mst_edges = []
|
||
|
|
|
||
|
|
while edge_heap and len(mst_edges) < self.V - 1:
|
||
|
|
weight, u, parent = heapq.heappop(edge_heap)
|
||
|
|
if mst_set[u]: # 忽略已在 MST 中的节点
|
||
|
|
continue
|
||
|
|
|
||
|
|
mst_set[u] = True
|
||
|
|
total_cost += weight
|
||
|
|
|
||
|
|
# 将加入的边存储（跳过起始节点的 -1 父节点）
|
||
|
|
if parent != -1:
|
||
|
|
mst_edges.append((parent, u, weight))
|
||
|
|
|
||
|
|
# 将相邻边加入优先队列
|
||
|
|
for v, w in self._adjacent_edges(u):
|
||
|
|
if not mst_set[v]:
|
||
|
|
heapq.heappush(edge_heap, (w, v, u))
|
||
|
|
|
||
|
|
return total_cost, mst_edges
|
||
|
|
|
||
|
|
def _adjacent_edges(self, u):
|
||
|
|
# 返回与顶点 u 相连的所有边
|
||
|
|
return [(v, weight) for (u_, v, weight) in self.graph if u_ == u] + \
|
||
|
|
[(u, weight) for (u, v_, weight) in self.graph if v_ == u]
|
||
|
|
|
||
|
|
# Kruskal 算法
|
||
|
|
def kruskal_mst(self):
|
||
|
|
self.graph.sort(key=lambda x: x[2]) # 按权值排序
|
||
|
|
parent = list(range(self.V))
|
||
|
|
rank = [0] * self.V
|
||
|
|
|
||
|
|
def find(u):
|
||
|
|
if parent[u] != u:
|
||
|
|
parent[u] = find(parent[u])
|
||
|
|
return parent[u]
|
||
|
|
|
||
|
|
def union(u, v):
|
||
|
|
root_u = find(u)
|
||
|
|
root_v = find(v)
|
||
|
|
if rank[root_u] > rank[root_v]:
|
||
|
|
parent[root_v] = root_u
|
||
|
|
elif rank[root_u] < rank[root_v]:
|
||
|
|
parent[root_u] = root_v
|
||
|
|
else:
|
||
|
|
parent[root_v] = root_u
|
||
|
|
rank[root_u] += 1
|
||
|
|
|
||
|
|
mst_edges = []
|
||
|
|
total_cost = 0
|
||
|
|
|
||
|
|
for u, v, weight in self.graph:
|
||
|
|
if find(u) != find(v):
|
||
|
|
union(u, v)
|
||
|
|
mst_edges.append((u, v, weight))
|
||
|
|
total_cost += weight
|
||
|
|
|
||
|
|
return total_cost, mst_edges
|