如何使用 Python 堆实现 Dijkstra 算法?

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

Dijkstra 算法是解决单源最短路径问题的经典算法,其基本思想是从源点开始,每次选取离源点最近的一个点,并更新其它点到源点的距离。Python 中可以使用 heapq 模块来实现堆,从而优化 Dijkstra 算法的时间复杂度。

以下是使用 Python 堆实现 Dijkstra 算法的代码示例,以图形表示的方式演示算法步骤。

import heapq

def dijkstra(graph, start):
    """通过堆实现 Dijkstra 算法"""
    distances = {node: float('inf') for node in graph}  # 初始化距离为无穷大
    distances[start] = 0  # 起点到自己的距离为0
    pq = [(0, start)]  # 创建一个优先队列,起点距离为0
    while pq:
        (dist, node) = heapq.heappop(pq)  # 获取距离最小的顶点
        if dist > distances[node]:  # 如果已经找到更优解,则跳过
            continue
        for neighbor, weight in graph[node].items():  # 遍历顶点的邻居
            distance = dist + weight
            if distance < distances[neighbor]:  # 更新邻居到起点的距离
                distances[neighbor] = distance
                heapq.heappush(pq, (distance, neighbor))  # 添加到堆中
    return distances

# 示例图
graph = {'A': {'B': 6, 'D': 1},
         'B': {'A': 6, 'C': 5, 'D': 2},
         'C': {'B': 5, 'D': 5, 'E': 3},
         'D': {'A': 1, 'B': 2, 'C': 5, 'E': 4},
         'E': {'C': 3, 'D': 4}}

# 从 A 点出发,求到其它点的最短距离
distances = dijkstra(graph, 'A')
print(distances)  # {'A': 0, 'B': 3, 'C': 8, 'D': 1, 'E': 7}

以上代码将图形转换成了邻接表的形式,其中每个顶点都是一个键值对,键为顶点名称,值为一个字典,该字典以邻居顶点为键,边权重为值。例如,graph['A'] 表示 A 点的邻居顶点有 B 和 D,对应的边权重为 6 和 1。

在 Dijkstra 算法的实现中,使用字典 distances 来保存起点到每个顶点的距离值。通过堆来维护已经访问过的顶点和它们到起点的距离值,每次从堆中取出距离最小的顶点,并遍历它的邻居。如果邻居的距离值可以通过当前顶点到该邻居的边权重更新,就将该邻居的距离和堆中标记为已访问。最后返回 distances 字典。

最终输出的结果为 {'A': 0, 'B': 3, 'C': 8, 'D': 1, 'E': 7},表示从 A 到其它点的最短距离。例如,A 到 B 的最短距离为 3,A 到 E 的最短距离为 7。

相关文章