1197_최소 스패닝 트리

입력

첫째 줄에는 정점의 개수(V), 간선의 개수(E)를 입력
두번째 줄부터 E+1번째 줄까지 (정점1, 정점2, 가중치)를 입력

출력

최소 스패닝 트리의 가중치를 출력

조건

• 정점 사이에 사이클이 존재할 수 있음

풀이

MST 문제를 푸는 방법으로는 크루스칼 알고리즘과 프림 알고리즘이 존재한다.
크루스칼 알고리즘은 Union Find 알고리즘을 사용하고 프림은 우선순위 큐를 사용하는 특징이 있다.
Union Find 알고리즘을 아직 완벽히 이해하고 있지 않아 프림 알고리즘으로 문제를 해결했다.

1. 정점마다 연결되는 정보를 확인하기 위한 딕셔너리(adj), 정점 방문 여부 리스트(check), 우선순위 큐(pq), 생성

adj = {v:[] for v in range(1, V+1)}
check = [False for _ in range(V)]
ans = 0
cnt = 0
pq = []

2. 입력받는 각 정점의 정보를 추가하고 정점 "1"의 정보로 우선순위 큐를 초기 설정

for _ in range(E):
    A, B, C = map(int, sys.stdin.readline().split())
    adj[A].append((C, B))
    adj[B].append((C, A))

# 우선순위 큐 초기 설정
check[0] = True
cnt += 1
for e in adj[1]: heapq.heappush(pq, (e[0], 1, e[1]))

3. 현재 정점과 연결된 정점 중 weight가 가장 작은 정점을 받으며 그 정점이 방문하지 않은 정점이라면 해당 weight 값을 ans에 추가하고 current 정점으로 설정
   방문한 정점의 개수가 V와 동일해진다고 종료

while pq:
    weight, curr, v = heapq.heappop(pq)

    if not check[v-1]:
        curr = v
        for e in adj[curr]: heapq.heappush(pq, (e[0], v, e[1]))

        check[curr-1] = True
        ans += weight
        cnt += 1
    else: continue

    if cnt == V: break

전체 코드

최소 스패닝 트리