最小生成树——Prim/Kruskal Python

最小生成树

从一个图中，生成一个权重最小的生成树

Prim

朴素版 $O(n^2)$ 稠密图

不断重复以下过程：

• 选择与当前集合距离最近的点，加入集合
• 拓展当前集合

和Dijsktra的思想类似，每次拓展与当前集合最近的点（而不再是原点）
对应的唯一区别就是

for j in range(1, n+1):
        dis[j] = min(dis[j], g[t][j])

在Prim里是直接取g[t][j]，而dijsktra是取dis[t]+g[t][j]

n,m = map(int, input().split())

g = [[float('inf') ] * (n+1) for _ in range(n+1)]

dis = [float('inf') ]* (n+1)

for i in range(m):
    x,y,z = map(int, input().split())
    g[x][y] = min(g[x][y], z)
    g[y][x] = g[x][y]

st = [False] *(n+1)

res = 0
for i in range(n):
    t = -1
    # 找到距离集合最近的点
    for j in range(1,n+1):
        if not st[j] and (t == -1 or dis[j] < dis[t]):
            t = j
    if i:
        res += dis[t]
    st[t] = True

    for j in range(1, n+1):
        dis[j] = min(dis[j], g[t][j])

if res == float('inf'):
    print('impossible')
else:
    print(res)

优化版 $O(m\log n)$

和迪杰斯特拉的优化方式一样，使用堆来优化，改成bfs，每次取出队列中距离集合最短的点拓展。

Kruskal 稀疏图

算法流程：

• 将所有的边按权重从小到大排序
• 每次取出权重最小的边，如果两个点不属于一个集合，融合他们（并查集）

import heapq


n,m = map(int, input().split())
# 首先读入所有的边，存入堆中
edges = []
for i in range(m):
    x,y,z = map(int, input().split())
    edges.append((z,x,y))

heapq.heapify(edges)
# 设置n个点的并查集
s = [i for i in range(n+1)]

def find(x):
    if s[x] != x:
        s[x] = find(s[x])
    return s[x]

# 开始出堆，每次取出堆顶元素，融合集合
res = 0
cnt = 0
while edges:
    t = heapq.heappop(edges)
    z,x,y = t
    if find(x) != find(y):
        s[find(x)] = find(y)
        res += z
        cnt += 1

if cnt == n-1:
    print(res)
else:
    print('impossible')