Cod sursa(job #3335826)

Utilizator efiruFiru Eduard efiru Data 23 ianuarie 2026 17:48:50
Problema Arbore partial de cost minim Scor 0
Compilator py Status done
Runda Arhiva educationala Marime 2.07 kb
# class DSU:
#     def __init__(self , n: int):
#         self.parent = list(range(n + 1))
#         self.size = [1 for _ in range(n + 1)]
#
#     def find(self , x: int) -> int:
#         if self.parent[x] != x:
#             self.parent[x] = self.find(self.parent[x])
#         return self.parent[x]
#
#     def union(self , x: int , y: int) -> bool:
#         rx , ry = self.find(x) , self.find(y)
#
#         if rx == ry:
#             return False
#
#         if self.size[rx] < self.size[ry]:
#             rx , ry = ry , rx
#
#         self.parent[ry] = rx
#         self.size[rx] += self.size[ry]
#         return True
#
#
#
#
with open("apm.in" , "r") as f:
    n , m = map(int , f.readline().split())

    edges = []
    adj = [[] for _ in range(n)]

    for _ in range(m):
        u , v , w = map(int , f.readline().split())
        edges.append((u , v , w))
        adj[u].append((w, v))

# edges.sort(key=lambda e : e[2])
# dsu = DSU(n)
# MST = []
# cost = 0
# taken = 0
#
# for u , v , w in edges:
#     if dsu.union(u , v):
#         cost += w
#         taken += 1
#         MST.append((u , v))
#
#         if taken == n - 1:
#             break
#
#

INF = 10 ** 18
import heapq

def prim():
    cost = 0
    taken = 0

    in_apm = [False * (n + 1)]
    key = [INF * (n + 1)]
    parent = [-1 * (n + 1)]

    heap = []
    heapq.heappush(heap , (0 , 1))

    while heap:
        w , u = heapq.heappop(heap)
        if in_apm[u]:
            continue
        
        in_apm[u] = True
        cost += w
        taken += 1

        for wv , v in adj[u]:
            if wv < key[v]:
                key[v] = wv
                parent[v] = u
                heapq.heappush(heap , (key[v] , v))
        
        APM = []
        
        start = 1
        for v in range(1, n + 1):
            if v != start and parent[v] != -1:
                APM.append((parent[v] , v , key[v]))

    return (APM , cost , taken)


    


MST , cost , taken = prim()

with open("apm.out" , "w") as g:
    g.write(str(cost) + "\n")
    g.write(str(taken) + "\n")

    for u , v in MST:
        g.write(f"{u} {v}\n")