0590 - Prim

Cerința

Se dă un graf neorientat ponderat conex cu n vârfuri și m muchii – în care fiecare muchie are asociat un cost, număr natural strict pozitiv. Folosind algoritmul lui Prim, determinați un arbore parțial de cost minim, cu rădăcina în vârful 1.

Date de intrare

Fișierul de intrare primIN.txt conține pe prima linie numerele n m, iar următoarele linii câte un triplet i j c, cu semnificația: există muchia (i j) și are costul c.

Date de ieșire

Fișierul de ieșire primOUT.txt va conține pe prima linie costul arborelui de cost minim determinat, iar pe a doua linie vectorul de tați al arborelui, cu elementele separate prin exact un spațiu. În cazul în care restricțiile nu sunt îndeplinite, se va afișa mesajul "Datele nu corespund restrictiilor impuse".

Restricții și precizări

• 1 ≤ n ≤ 100
• costul unei muchii va fi mai mic decât 1000

Exemplul 1:

primIN.txt

7 11
1 2 2
1 7 4
2 3 3
2 5 2
2 6 3
2 7 3
3 4 1
3 5 2
4 5 1
5 6 3
6 7 5

primOUT.txt

12
0 1 4 5 2 2 2 

Exemplul 2:

primIN.txt

0 11
1 2 2
1 7 4
2 3 3
2 5 2
2 6 3
2 7 3
3 4 1
3 5 2
4 5 1
5 6 3
6 7 5

primOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

<syntaxhighlight lang="python" line="1"> import heapq

def check_restrictions(n, edges):

   if not (1 <= n <= 100):
       return False
   for x, y, c in edges:
       if c >= 1000:
           return False
   return True

def main():

   with open("primIN.txt", "r") as infile:
       n, m = map(int, infile.readline().strip().split())
       edges = []
       for _ in range(m):
           x, y, c = map(int, infile.readline().strip().split())
           edges.append((x, y, c))

   if not check_restrictions(n, edges):
       with open("primOUT.txt", "w") as outfile:
           outfile.write("Datele nu corespund restrictiilor impuse")
       return

   G = [[] for _ in range(n + 1)]
   for x, y, c in edges:
       G[x].append((c, y))
       G[y].append((c, x))

   V = [False] * (n + 1)
   T = [-1] * (n + 1)
   D = [float('inf')] * (n + 1)

   Q = []
   V[1] = True
   T[1] = 0
   D[1] = 0

   for c, y in G[1]:
       T[y] = 1
       D[y] = c
       heapq.heappush(Q, (c, y))

   S = 0
   for k in range(1, n):
       while V[Q[0][1]]:
           heapq.heappop(Q)
       c, u = heapq.heappop(Q)
       V[u] = True
       S += c
       for cost, v in G[u]:
           if not V[v] and D[v] > cost:
               T[v] = u
               D[v] = cost
               heapq.heappush(Q, (cost, v))

   with open("primOUT.txt", "w") as outfile:
       outfile.write(f"{S}\n")
       outfile.write(" ".join(map(str, T[1:])))

if __name__ == "__main__":

   main()

</syntaxhighlight>