Algoritmia - codigo de primm

import heapq

class Grafo:

    def __init__(self, vertices):

        self.vertices = vertices

        self.grafo = [[] for _ in range(vertices)]

    def agregar_arista(self, u, v, peso):

        self.grafo[u].append((v, peso))

        self.grafo[v].append((u, peso))

    def prim(self):

        MST = []

        visitado = [False] * self.vertices

        heap = [(0, 0)]  # (peso, vértice)

        while heap:

            peso, actual = heapq.heappop(heap)

            if not visitado[actual]:

                visitado[actual] = True

                MST.append((peso, actual))

                for vecino, peso_vecino in self.grafo[actual]:

                    if not visitado[vecino]:

                        heapq.heappush(heap, (peso_vecino, vecino))

        return MST

# Ejemplo de uso

g = Grafo(5)

g.agregar_arista(0, 1, 2)

g.agregar_arista(0, 3, 8)

g.agregar_arista(1, 2, 3)

g.agregar_arista(1, 3, 5)

g.agregar_arista(1, 4, 1)

g.agregar_arista(2, 4, 7)

g.agregar_arista(3, 4, 4)

resultado_prim = g.prim()

print("Árbol de expansión mínima (MST) obtenido por el algoritmo de Prim:")

for peso, vertice in resultado_prim:

    print(f"Arista: {vertice} - Peso: {peso}")