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

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()