Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #!/usr/bin/env python3
- import sys
- from utils import *
- from collections import defaultdict
- from heapq import *
- file = open(sys.argv[1], "r")
- line = [int(x) for x in file.readline().split()]
- n = line[0] # vertices
- m = line[1] # edges
- position_of_extra_node = n
- n = n + 1 # extra S node
- edges = [[] for i in range(n)]
- edge_costs = [[0 for j in range(n)] for i in range(n)]
- overwritten_at = [0] * n
- vertex_costs = [float("inf")] * n
- for i in range(0, m):
- line = [int(x) for x in file.readline().split()]
- edges[line[0]].append(line[1])
- edges[line[1]].append(line[0])
- edge_costs[line[0]][line[1]] = line[2]
- edge_costs[line[1]][line[0]] = line[2]
- #print(edges)
- #print("costs")
- #print(edge_costs)
- is_terminals = [False] * n
- first_terminal = -1
- number_of_terminals = 0
- terminals = []
- for x in file.readline().split():
- if first_terminal == -1:
- first_terminal = int(x)
- terminals.append(int(x))
- is_terminals[int(x)] = True
- number_of_terminals = number_of_terminals + 1
- file.close()
- # first 0 edge
- edges[position_of_extra_node].append(first_terminal)
- edges[first_terminal].append(position_of_extra_node)
- vertex_costs[position_of_extra_node] = 0
- overwritten_at[position_of_extra_node] = 1
- #print("the source")
- #print(position_of_extra_node)
- previous = [0] * n
- def backtrack_path_and_assign_zero(node):
- #print("==== === BACKTRACKING")
- #print(previous)
- if node == first_terminal:
- vertex_costs[position_of_extra_node] = 0
- vertex_costs[node] = 0
- return
- while vertex_costs[node] != 0:
- #print("Add EDGE from: ")
- # print(node)
- vertex_costs[node] = 0
- edges[node].append(position_of_extra_node)
- edges[position_of_extra_node].append(node)
- node = previous[node]
- def _print(node, cost, queue):
- print("====== NEXT ===== ")
- print("Node: ")
- print(node)
- print("Of Cost:")
- print(cost)
- print("Remaining in Q:")
- print(len(queue))
- def _print1(node, visited, considered_cost):
- print("- Visited:")
- print(visited[node])
- print("- New Cost:")
- print(considered_cost)
- print("_ Old Cost:")
- print(vertex_costs[node])
- def dijkstra(run):
- queue = []
- visited = [False] * n
- visited[position_of_extra_node] = True
- for x in edges[position_of_extra_node]:
- heappush(queue, (0, x))
- while len(queue) != 0:
- (cost, node) = heappop(queue)
- # _print(node, cost, queue) # debug
- if is_terminals[node] and vertex_costs[node] != 0:
- # print("==== TERMINATING")
- backtrack_path_and_assign_zero(node)
- return
- if visited[node]:
- # print("=CANT HELP HERE")
- continue
- if cost > vertex_costs[node]:
- continue
- visited[node] = True
- # print("== NEIGHBORROWS:")
- for x in edges[node]:
- # print("Node: ")
- # print(x)
- considered_cost = edge_costs[node][x] + vertex_costs[node]
- # _print1(x, visited, considered_cost)
- if considered_cost < vertex_costs[x]:
- # print("-- ADD IT!")
- overwritten_at[x] = run
- vertex_costs[x] = considered_cost
- previous[x] = node
- heappush(queue, (vertex_costs[x], x))
- elif not visited[x] and overwritten_at[x] < run and vertex_costs[x] == considered_cost:
- # print("-- ADD IT anyway!")
- heappush(queue, (vertex_costs[x], x))
- for i in range(number_of_terminals):
- # print("================= NEW ROUND")
- # adjusted_dijkstra()
- dijkstra(i + 1)
- result_edges = []
- def backtrack_result(node):
- cost = 0
- while edge_costs[node][previous[node]] != 0:
- result_edges.append((node, previous[node]))
- cost = cost + edge_costs[node][previous[node]]
- edge_costs[node][previous[node]] = 0
- edge_costs[previous[node]][node] = 0
- node = previous[node]
- return cost
- result_cost = 0
- for x in terminals:
- result_cost = result_cost + backtrack_result(x)
- #print("===== THE RESULT =====")
- #print(result_cost)
- #print("===== THE EDGES =====")
- #for x in result_edges:
- # print(x)
- # print("last terminal ")
- # print(last_terminal)
- # print(terminals)
- file1 = open(sys.argv[2], "w")
- file1.write(str(result_cost))
- file1.write("\n")
- for i in result_edges:
- file1.write(str(i[0]))
- file1.write(" ")
- file1.write(str(i[1]))
- file1.write("\n")
- file1.close()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement