Python中的Dijkstra算法实现

2023-04-17 00:00:00 python dijkstra 算法

Dijkstra算法是一种最短路径算法，用于在加权图中找到从源节点到目标节点的最短路径。下面是Python中的Dijkstra算法实现代码：

import heapq

def dijkstra(graph, start):
    # 初始化距离和前驱字典
    distances = {node: float('inf') for node in graph}
    predecessors = {node: None for node in graph}
    # 堆中存储[start, 0]，表示源节点到自身的距离为0
    heap = [(start, 0)]
    # 源节点到自身的距离为0
    distances[start] = 0

    while heap:
        # 弹出堆顶元素
        current_node, distance = heapq.heappop(heap)
        # 如果当前节点已经被遍历过了，则跳过
        if distance > distances[current_node]:
            continue
        for neighbor, weight in graph[current_node].items():
            # 计算源节点到当前邻居节点的距离
            new_distance = distance + weight
            # 如果新距离更短，则更新距离和前驱字典，并将邻居节点入堆
            if new_distance < distances[neighbor]:
                distances[neighbor] = new_distance
                predecessors[neighbor] = current_node
                heapq.heappush(heap, (neighbor, new_distance))
    return distances, predecessors

这个实现中，graph是表示加权图的字典，以每个节点作为键，以与该节点相邻节点及其权值的字典作为值。例如，一个由a、b、c三个节点组成，其中a到b的权重为2，b到c的权重为3，a到c无权重的图可以表示为：

graph = {
    'a': {'b': 2},
    'b': {'c': 3},
    'c': {}
}

start是表示起点的字符串。

函数返回一个包含最短距离和前驱节点的元组。其中，distances是一个字典，以每个节点作为键，以从源节点到达该节点的最短距离作为值。例如，从a到每个节点的最短距离为：

{'a': 0, 'b': 2, 'c': 5}

predecessors是一个字典，以每个节点作为键，以源节点到达该节点的最短路径上的前驱节点作为值。例如，从a到每个节点的前驱节点为：

{'a': None, 'b': 'a', 'c': 'b'}

使用起来，只需要传递图和起点即可，例如，使用上述图计算从a到每个节点的最短距离和前驱节点：

distances, predecessors = dijkstra(graph, 'a')
print(distances)        # {'a': 0, 'b': 2, 'c': 5}
print(predecessors)     # {'a': None, 'b': 'a', 'c': 'b'}

注意，上述实现中，堆中存储的是包含节点和距离的元组，以距离作为排序依据，这样可以保证每次弹出的是距离最短的节点。如果两个节点到源节点的距离相等，则按节点名称的字典序排序。此外，为了优化算法性能，使用了前驱节点字典预处理最短路径。

相关文章