Python实现dijkstra算法

dijkstra.py:

import yaml
import copy

class Dijkstra:
    def __init__(self, path):
        # 读取路径配置
        self.data = self._config_reader(path)
        # 初始化最短路径信息
        self.path_data_init()

    def _config_reader(self, path):
        with open(path, 'r', encoding='utf8') as f:
            return yaml.load(f, yaml.BaseLoader)

    def path_data_init(self):
        data = {}
        for n in self.data.keys():
            data.update({
                n: {
                    'cost': float('inf'),
                    'path': []
                }
            })
        self.path_data = data
    
    def get_shortest_path(self, node_data, visited_path=[], base_cost=0):
        for node, cost in node_data.items():
            cost = int(cost)
            # 如果这个节点已经走过了，则跳过循环
            if node in visited_path:
                continue
            # 必须使用深拷贝，记录已经走过的路径
            visited_path_copy = copy.deepcopy(visited_path)
            visited_path_copy.append(node)
            # 如果总计的 Cost 值小于 path_data 记录的 Cost 值，
            # 则认为是最短路径，并更新到 path_data 中
            if base_cost + cost < self.path_data[node]['cost']:
                self.path_data[node]['cost'] = base_cost + cost
                self.path_data[node]['path'] = visited_path_copy
                # 递归执行，node_data 的值为 self.data[node]，表示走到下个节点
                self.get_shortest_path(self.data[node], visited_path_copy, base_cost + cost)
    
    def run(self, start_node):
        # 初始化起始节点的 cost 和 path 值
        self.path_data[start_node]['cost'] = 0
        self.path_data[start_node]['path'] = [start_node]
        # 获取最短路径
        self.get_shortest_path(self.data[start_node], visited_path=[start_node])

if __name__ == '__main__':
    dj = Dijkstra('dj.yaml')
    # 获取从 s0 节点为起始点的最短路径
    dj.run('s0')
    # 打印最短路径信息
    print(dj.path_data)
    # 清除最短路径信息
    dj.path_data_init()
    # 再次获取 s1 节点的最短路径
    dj.run('s1')
    print(dj.path_data)

dj.yaml:

s0:
  s1: 2
  s2: 6

s1:
  s0: 2
  s3: 5

s2:
  s0: 6
  s3: 8

s3:
  s1: 5
  s2: 8
  s4: 10
  s5: 15

s4:
  s3: 10
  s5: 6
  s6: 2

s5:
  s3: 15
  s4: 6
  s6: 6

s6:
  s4: 2
  s5: 6
  s7: 1

s7:
  s6: 1