Python实现最小生成树-Prim和Kruskal算法

1. Prim算法

Prim算法用于找到一个连通加权无向图的最小生成树。它的基本思想是从一个节点开始，每次选择边权最小且不会形成环的一条边，直到包括图中所有的节点。

1.1 算法步骤：

1. 从任意节点开始，初始化一个节点集合，将其标记为已访问。
2. 找到所有连接已访问节点和未访问节点的边，选择权重最小的边，并将其连接的未访问节点加入集合。
3. 重复步骤2，直到所有节点都被访问。

1.2 Python实现Prim算法

import heapq

def prim(graph):
    # 初始化最小生成树和总权重
    mst = []
    total_weight = 0
    # 从第一个节点开始
    start_node = list(graph.keys())[0]
    # 使用优先队列存储边
    edges = [(0, start_node, start_node)]
    visited = set()
    
    while edges:
        weight, u, v = heapq.heappop(edges)
        if v in visited:
            continue
        visited.add(v)
        mst.append((u, v, weight))
        total_weight += weight
        
        for next_node, next_weight in graph[v].items():
            if next_node not in visited:
                heapq.heappush(edges, (next_weight, v, next_node))
    
    return mst, total_weight

# 示例图表示为邻接表
graph = {
    'A': {'B': 1, 'C': 4},
    'B': {'A': 1, 'C': 2, 'D': 5},
    'C': {'A': 4, 'B': 2, 'D': 1},
    'D': {'B': 5, 'C': 1}
}

mst, total_weight = prim(graph)
print("最小生成树:", mst)
print("总权重:", total_weight)

2. Kruskal算法

Kruskal算法用于找到一个连通加权无向图的最小生成树。它的基本思想是每次选择权重最小的边，如果这条边的加入不会形成环，则将其加入最小生成树中，直到包括图中所有的节点。

2.1 算法步骤：

1. 将所有边按权重从小到大排序。
2. 初始化一个并查集（Union-Find）来记录每个节点所属的集合。
3. 按权重从小到大选择边，如果边连接的两个节点不属于同一个集合，则将其加入最小生成树，并合并这两个节点的集合。
4. 重复步骤3，直到最小生成树包含所有节点。

2.2 Python实现Kruskal算法

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
    
    def find(self, u):
        if self.parent[u] != u:
            self.parent[u] = self.find(self.parent[u])
        return self.parent[u]
    
    def union(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)
        if root_u != root_v:
            if self.rank[root_u] > self.rank[root_v]:
                self.parent[root_v] = root_u
            elif self.rank[root_u] < self.rank[root_v]:
                self.parent[root_u] = root_v
            else:
                self.parent[root_v] = root_u
                self.rank[root_u] += 1

def kruskal(graph):
    # 将边按权重排序
    edges = sorted(graph['edges'], key=lambda edge: edge[2])
    # 初始化并查集
    uf = UnionFind(graph['num_nodes'])
    mst = []
    total_weight = 0
    
    for u, v, weight in edges:
        if uf.find(u) != uf.find(v):
            uf.union(u, v)
            mst.append((u, v, weight))
            total_weight += weight
    
    return mst, total_weight

# 示例图表示为边列表
graph = {
    'num_nodes': 4,
    'edges': [
        (0, 1, 1),
        (0, 2, 4),
        (1, 2, 2),
        (1, 3, 5),
        (2, 3, 1)
    ]
}

mst, total_weight = kruskal(graph)
print("最小生成树:", mst)
print("总权重:", total_weight)

这两个算法分别解决了找到连通加权无向图的最小生成树的问题。Prim算法适用于稠密图，因为它主要操作节点；Kruskal算法则适用于稀疏图，因为它主要操作边。