Cod sursa(job #3196550)

Utilizator NCT127Cristina NCT127 Data 24 ianuarie 2024 11:07:39
Problema Arbore partial de cost minim Scor 70
Compilator py Status done
Runda Arhiva educationala Marime 1.65 kb
import heapq
f = open("apm.in")
# N reprezinta nr de noduri, M reprezinta nr de muchii
N,M = [int(x) for x in f.readline().split()]


def lista_adiacenta(nr_noduri, nr_muchii,orientare):
    graf = dict()
    for i in range(1,nr_noduri+1):
        graf[i] = []
    
    for i in range(nr_muchii):
        x,y,w = [int(u) for u in f.readline().split()]
        if orientare:
            graf[x].append((y,w))
        else:
            graf[x].append((y,w))
            graf[y].append((x, w))
    return graf

start=1
# prim
inf = float('inf')
heap = []
tata = [0]*(N+1)
viz  = [0]*(N+1)
heapq.heapify(heap)
heapq.heappush(heap,(0,start))
dis = [inf]*(N+1)
dis[start] = 0
graf=lista_adiacenta(N,M,0)

while heap:
    while heap and viz[heap[0][1]]:
        heapq.heappop(heap)
    
    if not heap:
        break
    elem = heapq.heappop(heap)
    nod = elem[1]

    viz[nod]=1

    for vecin in graf[nod]:
        if viz[vecin[0]]==0 and dis[vecin[0]]>vecin[1]:
            dis[vecin[0]]=vecin[1]
            tata[vecin[0]] = nod
            heapq.heappush(heap,(dis[vecin[0]], vecin[0]))

# def find_cost(x, c):
#     for (y,w) in graf[x]:
#         if y==c:
#             return w
#     return -1

with open("apm.out","w") as g:
    cost =sum(dis[1:])
    g.write(str(cost)+'\n')
    numar_muchii_arbore = sum(1 for nod in tata[1:] if nod != 0)
    g.write(str(numar_muchii_arbore)+'\n')
    # cost_total=0
    # afisare arbore
    for nod in range(1,N+1):
        if tata[nod]!=0:
            # cost=find_cost(nod,tata[nod])
            # cost_total+=cost
            g.write(str(nod)+' '+str(tata[nod])+'\n')