python 单源最短路径

'''
带权有向图 G=(V,E) 每条边的权为非负数
给定V中的一个顶点，称为源
计算从源到其他各顶点的最短路长度
'''
import numpy as np
def Dijkstra(n, v, dist, prev, c):
    #s 集合包含已知的最短路径的顶点
    s = np.zeros(n, bool)
    #dist 记录从源到其他顶点的最短路径长度 prev 记录相应的最短路径 c[i][j]表示边(i,j)的权
    s[v] = True
    for i in range(n-1):
        temp = float('inf')
        u = v
        for j in range(n):
            if not s[j] and dist[j] < temp and dist[j] != 0:
                u = j
                temp = dist[j]
        s[u] = True
        for j in range(n):
            if not s[j] and c[u][j] != 0:
                newdist = dist[u] + c[u][j]
                if  newdist< dist[j] or dist[j] == 0:
                    dist[j] = newdist
                    prev[j] = u
    # 输出结果
    for i in range(1, n):
        if dist[i] == 0:
            print("inf: 1", end="")
        else:
            print('{}: 1'.format(dist[i]), end="")
        # 递归函数，因为prev中是从后往前读取节点
        Result(prev, i)
        print("->{}".format(i + 1))

#打印路径
def Result(prev, t):
    if prev[t] != 0:
        Result(prev, prev[t])
        print('->{}'.format(prev[t] + 1), end="")

n = int(input())
c = []
for i in range(n):
    c.append(list(map(int, input().rstrip().split())))
dist = np.array(c[0])
prev = np.zeros(n, int)
Dijkstra(n, 0, dist, prev, c)

测试样例：

输入：

5
0 10 0 30 100
0 0 50 0 0
0 0 0 0 10
0 0 20 0 60
0 0 0 0 0

输出：
10: 1->2
50: 1->4->3
30: 1->4
60: 1->4->3->5