[백준] 1197 최소 스패닝 트리 파이썬

문제

https://www.acmicpc.net/problem/1197

1197번: 최소 스패닝 트리

문제

그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

 

---

풀이

제목에서 대놓고 최소 스패닝 트리라고 주어졌다.

MST 문제는 크루스칼 또는 프림으로 풀어야 한다.

https://cme10575.tistory.com/115

그래프 알고리즘 정리 - 파이썬

두 알고리즘 모두 음의 간선이 포함되어도 된다.

크루스칼은 간선 기준, 프림은 정점 기준이기 때문에 간선의 개수가 적으면 크루스칼, 정점의 개수가 적으면 프림 알고리즘을 적용하는 편이 효율적이라 한다.

 

나는 프림으로 플었다.

프림은 방문한 정점들에 연결된 간선 중 가중치가 가장 작은 간선부터 연결시키는 알고리즘이다.

import sys
import heapq
input = sys.stdin.readline

v, e = map(int, input().split())
g = [[] for _ in range(v+1)]

for _ in range(e):
    a, b, w = map(int, input().split())
    g[a].append((w, b))
    g[b].append((w, a))


q = g[1] #1번 정점부터 시작
visited = [True, True] + [False]*(v-1)
heapq.heapify(q)
answer = 0
while q:
    w, dest = heapq.heappop(q)
    if not visited[dest]:
        visited[dest] = True
        answer += w
        for e in g[dest]:
            if not visited[e[1]]:
                heapq.heappush(q, e) #방문하지 않은 정점으로 연결되는 간선 추가

print(answer)