 # Untitled

Sep 25th, 2021 (edited)
226
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. from sys import stdin, stdout
2. import math
3. import heapq
4.
5. R = lambda : stdin.readline().strip()
6. RL = lambda f=None: list(map(f, R().split(' '))) if f else list(R().split(' '))
7.
8. output = lambda x: stdout.write(str(x) + '\n')
9. output_list = lambda x: output(' '.join(map(str, x)))
10.
11. MAX = int(5e5 + 5)
12.
13.
14.
15.
16. def djkstra(n, m, k, graph, cost, vis, a, b):
17.
18.     heap = []
19.     heapq.heappush(heap, [0, [1,0]])
20.     cost=0
21.     while len(heap):
22.         curr = heapq.heappop(heap)
23.         curr_n, curr_s = curr, curr
24.         cst = cost[curr_n][curr_s]
25.
26.         for i in range(len(graph[curr_n])):
27.             if graph[curr_n][i]==curr_n:
28.                 continue
29.             c_cst = cst + graph[curr_n][i]
30.             if (cost[graph[curr_n][i]][curr_s] == -1):
31.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i], curr_s]])
32.                 cost[graph[curr_n][i]][curr_s] = c_cst
33.             elif cost[graph[curr_n][i]][curr_s] > c_cst :
34.                 cost[graph[curr_n][i]][curr_s] = c_cst
35.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i], curr_s]])
36.
37.             c_cst = cst
38.             if curr_s == k:
39.                 continue
40.             if cost[graph[curr_n][i]][curr_s+1] == -1:
41.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i], curr_s+1]])
42.                 cost[graph[curr_n][i]][curr_s+1]=c_cst
43.             elif cost[graph[curr_n][i]][curr_s+1]>c_cst:
44.                 cost[graph[curr_n][i]][curr_s+1]=c_cst
45.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i], curr_s+1]])
46.
47.     ret = int(2e18)
48.     for i in range(1,n+1):
49.         for j in range(k+1):
50.             if cost[i][j] == -1:
51.                 continue
52.             ret = min(ret, cost[i][j])
53.         print(ret, i, j)
54.     return ret
55.
56.
57. n = int(R())
58. a = int(R())
59. b = int(R())
60. k = int(R())
61. m, _ = RL(int)
62.
63. graph = [ [] for i in range(n+1) ]
64. cost = [ 20*[-1] for i in range(n+1) ]
65. vis = [ 20*[False] for i in range(n+1) ]
66.
67. for i in range(m):
68.     u, v, w = RL(int)
69.     graph[u].append([v, w])
70.     graph[v].append([u, w])
71.
72. print(djkstra(n, m, k, graph, cost, vis, a, b))
73.
RAW Paste Data