Pagini recente » Cod sursa (job #190765) | Cod sursa (job #227721) | Cod sursa (job #2149054) | Cod sursa (job #592874) | Cod sursa (job #2936331)
# Stoica Ioan
# ex 2
# Path: apm_prim.py
# https://www.infoarena.ro/problema/apm
# using Prim's algorithm
# read from apm.in n,m and a list with edges and costs
with open('apm.in', 'r') as f:
n, m = [int(x) for x in f.readline().split(' ')]
edges = []
for line in f:
edges.append([int(x) for x in line.split()])
# create adjacency list
adj = [[] for i in range(n + 1)]
for edge in edges:
adj[edge[0]].append([edge[1], edge[2]])
adj[edge[1]].append([edge[0], edge[2]])
# visited nodes
visited = [False for i in range(n + 1)]
# initialize the priority queue
pq = []
# add the first node
visited[1] = True
for edge in adj[1]:
pq.append([1, edge[0], edge[1]])
# sort the priority queue
pq.sort(key=lambda x: x[2])
solution = []
cost = 0
# Prim's algorithm
while len(pq) > 0:
# get the edge with the smallest cost
edge = pq.pop(0)
if not visited[edge[1]]:
visited[edge[1]] = True
cost += edge[2]
solution.append(edge)
for e in adj[edge[1]]:
if not visited[e[0]]:
pq.append([edge[1], e[0], e[1]])
pq.sort(key=lambda x: x[2])
# write to apm.out the cost
with open('apm.out', 'w') as f:
f.write(str(cost) + '\n' + str(n-1) + '\n')
for edge in solution:
f.write(str(edge[0]) + ' ' + str(edge[1]) + '\n')