Cod sursa(job #2936331)

Utilizator IoanStoicaStoica Ioan IoanStoica Data 8 noiembrie 2022 17:28:02
Problema Arbore partial de cost minim Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 1.35 kb
# 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')