使用warshell算法就算最短路径,然后将据点的距离提出来进行排序,但是运行超时,不知道有没有大神指导一下?
import math
def shortDistance(dis):
node_num = len(dis)
for i in range(node_num): # 十字交叉法的位置位置,先列后行
for j in range(node_num): # 列 表示dis[j][i]的值,即j->i
for k in range(node_num): # 行 表示dis[i][k]的值,即i->k,i只是一个桥梁而已
# 先列后行,形成一个传递关系,若比原来距离小,则更新
if dis[j][k] > dis[j][i] + dis[i][k]:
dis[j][k] = dis[j][i] + dis[i][k]
return dis
n,m,k=list(map(int,input().strip().split()))
tag = list(map(int,input().strip().split()))
# 邻接矩阵初始化
graph = []
for i in range(n):
graph.append([math.inf]*n)
graph[i][i] = 0
for i in range(m):
temp = list(map(int,input().strip().split()))
temp[0] -= 1
temp[1] -= 1
if graph[temp[0]][temp[1]] > temp[2]:
graph[temp[0]][temp[1]]=temp[2]
graph[temp[1]][temp[0]]=temp[2]
dis = shortDistance(graph)
result = []
for i in range(n):
data = [] # 获取i点到每个据点的距离
for j in range(n):
if tag[j] == 1:
data.append(dis[i][j])
data.sort() # 排序
ans = 0 # 获取前K个最短距离之和
for j in range(k):
if data[j]!=math.inf:
ans += data[j]
else:
break
result.append(ans)
print("\n".join(map(str,result)))