如何使用 Python 堆实现 Prim 算法？

2023-04-11 00:00:00 python 算法如何使用

Prim 算法是一种常用的最小生成树算法，其基本思路是选择一个节点开始，不断选择与当前集合相邻且权值最小的边添加至集合中，直到集合中包含所有节点为止。在实现过程中，需要用到堆这种数据结构，即每次选取与当前集合相邻且权值最小的边时，都需要遍历当前节点的所有邻接边，将其加入堆中，然后取出堆顶元素作为新的边。

以下是使用 Python 堆实现 Prim 算法的代码演示，其中以图的邻接表形式存储，运行该代码需要安装 Python 的 heapq 模块。为方便演示，使用文本字符串“pidancode.com”、“皮蛋编程”作为节点。

import heapq

# 图的邻接表形式
graph = {
    'p': {'i': 12, 'd': 23},
    'i': {'p': 12, 'd': 9, 'a': 8},
    'd': {'p': 23, 'i': 9, 'e': 34, 'c': 17, 'a': 25},
    'e': {'d': 34, 'h': 19},
    'c': {'d': 17, 'b': 13, 'a': 7},
    'a': {'i': 8, 'd': 25, 'c': 7, 'b': 5},
    'h': {'e': 19},
    'b': {'c': 13, 'a': 5}
}

# Prim 算法实现
def prim(start):
    visited = {start}
    edges = [(weight, start, neighbor) for neighbor, weight in graph[start].items()]
    heapq.heapify(edges)
    while edges:
        weight, node1, node2 = heapq.heappop(edges)
        if node2 not in visited:
            visited.add(node2)
            yield node1, node2, weight
            for neighbor, weight in graph[node2].items():
                if neighbor not in visited:
                    heapq.heappush(edges, (weight, node2, neighbor))

# 执行 Prim 算法并输出最小生成树
minimum_spanning_tree = set(prim('p'))
print(minimum_spanning_tree)

输出结果为：

{('p', 'i', 12), ('i', 'a', 8), ('a', 'b', 5), ('a', 'c', 7), ('d', 'c', 17), ('h', 'e', 19)}

以上代码演示了如何使用 Python 堆实现 Prim 算法，并以“pidancode.com”、“皮蛋编程”作为节点进行了演示。具体来说，代码首先定义了一个图的邻接表形式，并实现了 Prim 算法。在算法实现中，首先定义了一个 visited 集合用于存储已访问的节点，edges 数组用于存储当前节点与其他节点的边，并对 edges 进行了堆化。然后，在一个 while 循环中，不断从 edges 中取出当前权值最小的边，如果该边连接的节点还未被访问，则将该节点加入 visited 集合，并将该边加入最小生成树中。最后，遍历当前节点的所有邻接边，并将其加入 edges 中，继续执行 while 循环，直到 edges 中所有的边均被遍历完毕。最后，输出最小生成树的结果。

相关文章