[Python]贪心算法-Dijkstra-实现

Posted 2020-12-12 Sight Tech.

目标

带权重的有向图上单源最短路径问题。且权重都为非负值。如果采用的实现方法合适，Dijkstra运行时间要低于Bellman-Ford算法。

思路

1. 选择一个节点开始蔓延
2. 计算自身到连接它的一级节点之间的距离， 全部作为候选集
3. 在候选集中，找到距离最短的，对应的那个节点
4. 删除这个节点在候选集中的信息
5. 继续蔓延，还是找最小的距离
6. 直到候选集为空

最小距离的判断标准 dist[j] = min(dist[j], dist[i] + weight[i][j])

完善版本

import heapq
import math

def dijkstra(graph, init_node):
    pqueue = []
    heapq.heappush(pqueue, (0, init_node))  # min heap, sort data item automatically
    visited = set()  # actually you dont have to use this.
    weight = dict.fromkeys(graph.keys(), math.inf)
    weight[init_node] = 0
    connection_dict = {init_node: "Path: Start From -> "}  # save connection records

    while len(pqueue) > 0:
        pair = heapq.heappop(pqueue)  # Pop the smallest item off the heap
        cost, start = pair[0], pair[1]
        visited.add(start)
        for end in graph[start].keys():
            if end not in visited and cost + graph[start][end] < weight[end]:
                # dist[j] = min(dist[j], dist[i] +  weight[i][j])
                heapq.heappush(pqueue, (cost + graph[start][end], end))
                connection_dict[end] = start
                weight[end] = cost + graph[start][end]

    return {v: k for k, v in connection_dict.items()}, weight


if __name__ == \'__main__\':
    graph_dict = {
        "A": {"B": 5, "C": 1},
        "B": {"A": 5, "C": 2, "D": 1},
        "C": {"A": 1, "B": 2, "D": 4, "E": 8},
        "D": {"B": 1, "C": 4, "E": 3, "F": 6},
        "E": {"C": 8, "D": 3},
        "F": {"D": 6},
    }

    path, distance = dijkstra(graph_dict, "A")
    print(path) # {\'Path: Start From -> \': \'A\', \'C\': \'B\', \'A\': \'C\', \'B\': \'D\', \'D\': \'F\'}
    print(distance) # {\'A\': 0, \'B\': 3, \'C\': 1, \'D\': 4, \'E\': 7, \'F\': 10}

最精简版本

import heapq

def dijkstra(graph, init_node):
    primary_queue = []
    heapq.heappush(primary_queue, (0, init_node))
    # the reason why i need to use this heap is because
    # i want to take advantage of its automatic sorting

    result = dict.fromkeys(graph.keys(), 123131)
    result[init_node] = 0

    while len(primary_queue) > 0:
        cost, start = heapq.heappop(primary_queue)

        for end in graph[start].keys():
            if result[start] + graph[start][end] < result[end]:
                # dist[j] = min(dist[j], dist[i] +  weight[i][j])
                heapq.heappush(primary_queue, (result[start] + graph[start][end], end))
                result[end] = result[start] + graph[start][end]

    return result

参考文章

1. 最短路知识点总结（Dijkstra，Floyd，SPFA，Bellman-Ford)
2. Python实现Dijkstra算法