import heapq

def prim(grafo):
    """
    Implementación del Algoritmo de Prim para encontrar el árbol generador mínimo.
    
    :param graph: Diccionario que representa el grafo ponderado conexo no dirigido,
                  donde las claves son los nodos y los valores son listas de tuplas (vecino, peso)
    :return: Lista de aristas que forman el árbol generador mínimo
    """
    # Inicializar el árbol generador mínimo
    arbol_gm = []
    
    # Seleccionar un nodo inicial arbitrario (por ejemplo, el primer nodo del grafo)
    nodo_inicial = list(grafo.keys())[0]
    
    # Crear una cola de prioridad para seleccionar la arista de menor peso
    bordes = [(peso, nodo_inicial, to) for to, peso in grafo[nodo_inicial]]
    heapq.heapify(bordes)
    
    # Conjunto de nodos visitados
    visitados = set([nodo_inicial])
    
    while bordes:
        # Seleccionar la arista de menor peso que conecta un nodo visitado con un nodo no visitado
        peso, frm, to = heapq.heappop(bordes)
        
        if to not in visitados:
            # Añadir la arista al árbol generador mínimo
            arbol_gm.append((frm, to, peso))
            visitados.add(to)
            
            # Añadir todas las aristas del nuevo nodo al heap
            for siguiente, peso in grafo[to]:
                if siguiente not in visitados:
                    heapq.heappush(bordes, (peso, to, siguiente))
    
    return arbol_gm

# Ejemplo 1 (pag 642  K. Rosen) / Solución Figura 2 (pag 643  K. Rosen)
ciudades = {
    'San Francisco': [('Denver', 900), ('Chicago', 1200), ('Nueva York', 2000), ('Atlanta', 2200)],
    'Denver': [('San Francisco', 900), ('Chicago', 1300), ('Nueva York', 1600), ('Atlanta', 1400)],
    'Chicago': [('San Francisco', 1200), ('Denver', 1300), ('Nueva York', 1000), ('Atlanta', 700)],
    'Nueva York': [('San Francisco', 2000), ('Denver', 1600), ('Chicago', 1000), ('Atlanta', 800)],
    'Atlanta': [('San Francisco', 2200), ('Denver', 1400), ('Chicago', 700), ('Nueva York', 800)]
}

# Ejemplo 2 (pag 642  K. Rosen) / Solución Figura 3 y 4 (pag 643  K. Rosen)
grafo = {
    'a': [('b', 2), ('e', 3)],
    'b': [('a', 2), ('c', 3), ('f', 1)],
    'c': [('b', 3), ('d', 1), ('g', 2)],
    'd': [('c', 1), ('h', 5)],
    'e': [('a', 3), ('f', 4), ('i', 4)],
    'f': [('b', 1), ('e', 4), ('g', 3), ('j', 2)],
    'g': [('c', 2), ('f', 3), ('h', 3), ('k', 4)],
    'h': [('d', 5), ('g', 3), ('l', 3)],
    'i': [('e', 4), ('j', 3)],
    'j': [('f', 2), ('i', 3), ('k', 3)],
    'k': [('g', 4), ('j', 3), ('l', 1)],
    'l': [('h', 3), ('k', 1)],
}

# Encontrar el árbol generador mínimo
agm_ciudades = prim(ciudades)
print("Árbol generador mínimo:", agm_ciudades)
agm = prim(grafo)
print("Árbol generador mínimo:", agm)