자료구조/👉알고리즘 - 26

from collections import defaultdict
from heapq import *

def prim(edges, start_node):
    tree = []
    # 각 노드별로 연결된 간선을 정리
    edgesOf = defaultdict(list)
    for weight, n1, n2 in edges:
        edgesOf[n1].append((weight, n1, n2))
        edgesOf[n2].append((weight, n2, n1))
    # 연결된 노드들
    connected_n = set(start_node)
    # 간선 리스트
    connected_e = edgesOf[start_node]
    heapify(connected_e)
    
    while connected_e:
        next_e = heappop(connected_e)
        if next_e[2] not in connected_n:
            # 트리에 경로 추가
            tree.append(next_e)
            # 연결된 노드에 추가
            connected_n.add(next_e[2])
            # 간선 리스트에 추가
            for edge in edgesOf[next_e[2]]:
                # 추가하려는 노드가 연결되지 않은 경우 추가
                if edge[2] not in connected_n:
                    heappush(connected_e, edge)

    return tree
    
    
graph = [
    (7, 'A', 'B'), (5, 'A', 'D'),
    (8, 'B', 'C'), (9, 'B', 'D'), (7, 'B', 'E'),
    (5, 'C', 'E'),
    (7, 'D', 'E'), (6, 'D', 'F'),
    (8, 'E', 'F'), (9, 'E', 'G'),
    (11, 'F', 'G')
]
prim(graph, 'A')