【算法】Python 实现最小生成树-Euler的博客

准备工作：图的表示

为了方便演示，我们假设图用以下方式表示：

顶点数量: 一个整数 V。
边列表 (Kruskal): 一个列表，每个元素是一个元组 (权重, 顶点1, 顶点2)。顶点用从 0 到 V-1 的整数表示。
邻接表 (Prim): 一个字典或列表的列表，graph[u] 包含所有与顶点 u 相邻的边，形式为 [(权重, 邻居顶点), ...]。

示例图

我们将使用以下加权无向图作为例子：

      (10)
   0--------1
   | \      |
(6)|  \(5)  |(15)
   |   \    |
   2--------3
      (4)

顶点: 0, 1, 2, 3 (V = 4)
边: (0, 1, 10), (0, 2, 6), (0, 3, 5), (1, 3, 15), (2, 3, 4)

1. Kruskal 算法实现

Kruskal 算法需要一个高效的方法来检测环，这通常通过**并查集（Disjoint Set Union, DSU）**数据结构实现。

# 并查集 (DSU) 数据结构实现
class DSU:
    def __init__(self, n):
        # 初始化每个节点的父节点为自身
        self.parent = list(range(n))
        # 可选：用于按秩合并优化
        self.rank = [0] * n

    def find(self, i):
        # 路径压缩：查找根节点，并将路径上所有节点的父节点直接指向根节点
        if self.parent[i] == i:
            return i
        self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, i, j):
        # 合并两个集合
        root_i = self.find(i)
        root_j = self.find(j)
        if root_i != root_j:
            # 按秩合并（可选优化）：将秩小的树合并到秩大的树上
            if self.rank[root_i] < self.rank[root_j]:
                self.parent[root_i] = root_j
            elif self.rank[root_i] > self.rank[root_j]:
                self.parent[root_j] = root_i
            else:
                self.parent[root_j] = root_i
                self.rank[root_i] += 1
            return True # 合并成功
        return False # 已经在同一集合，合并失败（表示会形成环）

# Kruskal 算法函数
def kruskal(num_vertices, edges):
    """
    使用 Kruskal 算法计算最小生成树。

    Args:
        num_vertices (int): 图中的顶点数量。
        edges (list): 边的列表，每个元素是 (权重, 顶点1, 顶点2)。

    Returns:
        tuple: (最小生成树的总权重, 构成最小生成树的边列表 [(u, v, w), ...])
               如果图不连通，可能返回不完整的树。
    """
    # 1. 按权重对所有边进行排序
    edges.sort()

    # 2. 初始化并查集
    dsu = DSU(num_vertices)

    mst_cost = 0
    mst_edges = []
    edges_count = 0 # 记录已加入 MST 的边数

    # 3. 遍历排序后的边
    for weight, u, v in edges:
        # 4. 检查加入此边是否会形成环 (find 操作)
        if dsu.union(u, v): # 如果 u 和 v 不在同一集合，则合并它们 (union 操作)
            # 5. 如果不形成环，则将此边加入 MST
            mst_cost += weight
            mst_edges.append((u, v, weight))
            edges_count += 1
            # 优化：如果已经找到 V-1 条边，MST 构建完成
            if edges_count == num_vertices - 1:
                break

    # 检查是否所有顶点都连通 (如果需要)
    if edges_count < num_vertices - 1 and num_vertices > 0:
        print("警告：图可能不连通，生成的不是包含所有顶点的生成树。")


    return mst_cost, mst_edges

# --- Kruskal 示例运行 ---
V_k = 4
edges_k = [
    (10, 0, 1),
    (6, 0, 2),
    (5, 0, 3),
    (15, 1, 3),
    (4, 2, 3)
]

mst_cost_k, mst_edges_k = kruskal(V_k, edges_k)
print("--- Kruskal 算法 ---")
print(f"最小生成树总权重: {mst_cost_k}")
print(f"最小生成树的边: {mst_edges_k}")
# 预期输出:
# --- Kruskal 算法 ---
# 最小生成树总权重: 15
# 最小生成树的边: [(2, 3, 4), (0, 3, 5), (0, 2, 6)]  (注意：(0,2,6) 和 (0,1,10) 取决于排序稳定性，但总权重和顶点覆盖不变，这里 (0,2,6) 权重更小)
# 修正预期：排序后是 (4, 2, 3), (5, 0, 3), (6, 0, 2), (10, 0, 1), (15, 1, 3)
# 选择 (4, 2, 3) -> union(2, 3) -> mst = [(2, 3, 4)], cost = 4
# 选择 (5, 0, 3) -> union(0, 3) -> mst = [(2, 3, 4), (0, 3, 5)], cost = 9
# 选择 (6, 0, 2) -> find(0) = find(2) ? 否 (根分别是0和0，但通过3相连，所以是同一个集合了)。find(0) 和 find(2) 的根现在相同了。跳过。
#   检查并查集状态：parent = [0, 1, 0, 0] after (2,3,4) and (0,3,5) unions. find(0)=0, find(2)=0. 是的，会形成环。
# 选择 (10, 0, 1) -> union(0, 1) -> mst = [(2, 3, 4), (0, 3, 5), (0, 1, 10)], cost = 19
# 修正：上面的 (6, 0, 2) 步骤错了。
# 重新来：
# 1. 边排序: [(4, 2, 3), (5, 0, 3), (6, 0, 2), (10, 0, 1), (15, 1, 3)]
# 2. 初始化 DSU: parent = [0, 1, 2, 3]
# 3. 选 (4, 2, 3): find(2)=2, find(3)=3. 不相等. union(2, 3). parent=[0, 1, 2, 2]. mst=[(2, 3, 4)], cost=4.
# 4. 选 (5, 0, 3): find(0)=0, find(3)=2. 不相等. union(0, 3). parent=[2, 1, 2, 2]. mst=[(2, 3, 4), (0, 3, 5)], cost=9.
# 5. 选 (6, 0, 2): find(0)=2, find(2)=2. 相等. 跳过.
# 6. 选 (10, 0, 1): find(0)=2, find(1)=1. 不相等. union(0, 1). parent=[2, 2, 2, 2]. mst=[(2, 3, 4), (0, 3, 5), (0, 1, 10)], cost=19.
# 7. 找到 V-1 = 3 条边，结束。
# 结果: cost=19, edges=[(2, 3, 4), (0, 3, 5), (0, 1, 10)]

# 再次检查示例图和权重... 啊，我手算的 MST 好像不对。
# 边: (0,1,10), (0,2,6), (0,3,5), (1,3,15), (2,3,4)
# Kruskal:
# 1. (2,3,4) - cost 4, mst={(2,3,4)}
# 2. (0,3,5) - cost 4+5=9, mst={(2,3,4), (0,3,5)}
# 3. (0,2,6) - 0和2已经通过3连接，跳过
# 4. (0,1,10) - cost 9+10=19, mst={(2,3,4), (0,3,5), (0,1,10)}
# 结束。总权重 19。

# Prim (从0开始):
# 1. 选0. 候选边: (0,2,6), (0,3,5), (0,1,10).
# 2. 选最小的 (0,3,5). cost=5. mst={(0,3,5)}. 访问 {0, 3}.
# 3. 候选边: (0,2,6), (0,1,10), (3,2,4), (3,1,15).
# 4. 选最小的 (3,2,4). cost=5+4=9. mst={(0,3,5), (3,2,4)}. 访问 {0, 3, 2}.
# 5. 候选边: (0,1,10), (3,1,15), (2,0,6) (已访问0).
# 6. 选最小的 (0,1,10). cost=9+10=19. mst={(0,3,5), (3,2,4), (0,1,10)}. 访问 {0, 3, 2, 1}.
# 结束。总权重 19。

# 看来我之前的预期输出写错了，正确的总权重是 19。代码逻辑是正确的。
print(f"修正后预期总权重: 19")
print(f"修正后预期边: [(2, 3, 4), (0, 3, 5), (0, 1, 10)]") # 边的顺序可能不同，但集合相同

2. Prim 算法实现

Prim 算法通常使用**优先队列（最小堆）**来高效地找到连接已选顶点和未选顶点的最小权重边。

import heapq # 导入优先队列库

# Prim 算法函数
def prim(num_vertices, graph_adj):
    """
    使用 Prim 算法计算最小生成树。

    Args:
        num_vertices (int): 图中的顶点数量。
        graph_adj (dict or list): 图的邻接表表示。
                                  例如: {0: [(1, 10), (2, 6), (3, 5)], ...}
                                  或 [[(1, 10), (2, 6), (3, 5)], ...]

    Returns:
        tuple: (最小生成树的总权重, 构成最小生成树的边列表 [(u, v, w), ...])
               如果图不连通，可能返回不完整的树。
    """
    if num_vertices == 0:
        return 0, []

    mst_cost = 0
    mst_edges = []
    visited = set()
    # 优先队列，存储 (权重, 起点, 终点)
    # 起点是已在 MST 中的顶点，终点是待加入的顶点
    min_heap = []

    # 1. 选择起始顶点 (例如顶点 0)
    start_node = 0
    visited.add(start_node)

    # 2. 将起始顶点的所有邻边加入优先队列
    # 检查 graph_adj 是字典还是列表
    if isinstance(graph_adj, dict):
        neighbors = graph_adj.get(start_node, [])
    elif isinstance(graph_adj, list) and start_node < len(graph_adj):
         neighbors = graph_adj[start_node]
    else:
        neighbors = [] # 或者抛出错误

    for weight, neighbor in neighbors:
        heapq.heappush(min_heap, (weight, start_node, neighbor))

    # 3. 主循环，直到访问了所有顶点或堆为空
    while min_heap and len(visited) < num_vertices:
        # 4. 从优先队列中取出权重最小的边
        weight, u, v = heapq.heappop(min_heap)

        # 5. 如果边的终点 v 已经访问过，则跳过 (避免形成环)
        if v in visited:
            continue

        # 6. 将顶点 v 加入已访问集合，并将边加入 MST
        visited.add(v)
        mst_cost += weight
        mst_edges.append((u, v, weight))

        # 7. 将 v 的所有连接到未访问顶点的邻边加入优先队列
        # 检查 graph_adj 是字典还是列表
        if isinstance(graph_adj, dict):
            v_neighbors = graph_adj.get(v, [])
        elif isinstance(graph_adj, list) and v < len(graph_adj):
             v_neighbors = graph_adj[v]
        else:
            v_neighbors = []

        for neighbor_weight, neighbor in v_neighbors:
            if neighbor not in visited:
                heapq.heappush(min_heap, (neighbor_weight, v, neighbor))

    # 检查是否所有顶点都连通
    if len(visited) < num_vertices and num_vertices > 0:
         print("警告：图可能不连通，生成的不是包含所有顶点的生成树。")

    return mst_cost, mst_edges

# --- Prim 示例运行 ---
V_p = 4
# 使用邻接表表示图
graph_adj_p = {
    0: [(10, 1), (6, 2), (5, 3)],  # 格式: (权重, 邻居)
    1: [(10, 0), (15, 3)],
    2: [(6, 0), (4, 3)],
    3: [(5, 0), (15, 1), (4, 2)]
}
# 或者使用列表形式（如果顶点是 0 到 V-1 连续整数）
# graph_adj_p_list = [
#     [(10, 1), (6, 2), (5, 3)], # 顶点 0 的邻居
#     [(10, 0), (15, 3)],       # 顶点 1 的邻居
#     [(6, 0), (4, 3)],         # 顶点 2 的邻居
#     [(5, 0), (15, 1), (4, 2)]  # 顶点 3 的邻居
# ]


mst_cost_p, mst_edges_p = prim(V_p, graph_adj_p)
print("\n--- Prim 算法 ---")
print(f"最小生成树总权重: {mst_cost_p}")
print(f"最小生成树的边: {mst_edges_p}")
# 预期输出:
# --- Prim 算法 ---
# 最小生成树总权重: 19
# 最小生成树的边: [(5, 0, 3), (4, 3, 2), (10, 0, 1)] (边的顺序和方向可能因实现而异，但集合和总权重相同)
print(f"预期总权重: 19")
print(f"预期边集合 (忽略顺序和方向): {{(0, 3, 5), (2, 3, 4), (0, 1, 10)}}")

代码说明

Kruskal:
- 核心在于 DSU 类，find 用于查找元素的根（代表），union 用于合并两个元素所在的集合。路径压缩和按秩合并是优化手段，确保操作接近常数时间复杂度。
- 算法首先对所有边按权重排序。
- 然后遍历排序后的边，如果一条边的两个顶点不在同一个集合中（dsu.union() 返回 True），则说明加入这条边不会形成环，将其加入 MST，并合并这两个顶点所在的集合。
Prim:
- 使用 heapq 模块实现最小优先队列，存储 (权重, 已在树中的顶点, 待加入顶点)。
- 从一个起始顶点开始，将其邻边加入堆。
- 循环地从堆中取出权重最小的边。如果这条边的目标顶点尚未访问，则将其标记为已访问，将边加入 MST，并将其所有连接到未访问顶点的邻边加入堆。
- 重复此过程，直到所有顶点都被访问。

目录CONTENT

【算法】Python 实现最小生成树

评论区