dark-Matter

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[1][0]=0
  21.     while len(heap):
  22.         curr = heapq.heappop(heap)
  23.         curr_n, curr_s = curr[1][0], curr[1][1]
  24.         cst = cost[curr_n][curr_s]
  25.  
  26.         for i in range(len(graph[curr_n])):
  27.             if graph[curr_n][i][0]==curr_n:
  28.                 continue
  29.             c_cst = cst + graph[curr_n][i][1]
  30.             if (cost[graph[curr_n][i][0]][curr_s] == -1):
  31.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i][0], curr_s]])
  32.                 cost[graph[curr_n][i][0]][curr_s] = c_cst
  33.             elif cost[graph[curr_n][i][0]][curr_s] > c_cst :
  34.                 cost[graph[curr_n][i][0]][curr_s] = c_cst
  35.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i][0], curr_s]])
  36.            
  37.             c_cst = cst
  38.             if curr_s == k:
  39.                 continue
  40.             if cost[graph[curr_n][i][0]][curr_s+1] == -1:
  41.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i][0], curr_s+1]])
  42.                 cost[graph[curr_n][i][0]][curr_s+1]=c_cst
  43.             elif cost[graph[curr_n][i][0]][curr_s+1]>c_cst:
  44.                 cost[graph[curr_n][i][0]][curr_s+1]=c_cst
  45.                 heapq.heappush(heap, [c_cst*-1, [graph[curr_n][i][0], 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