[백준] 1197번 : 최소 스패닝 트리 (python 파이썬)

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

 

import sys
import heapq
v,e = map(int, input().split(" "))

graph = [[] for _ in range(v + 1)]
for _ in range(e):
    start, end, score = map(int, sys.stdin.readline().split())
    graph[start].append([score, end])
    graph[end].append([score, start])


visit = [-1 for i in range(v+1)]
h = [[0, 1]]
answer = 0
while h:
    next = heapq.heappop(h)
    if visit[next[1]] != -1:
        continue
    visit[next[1]] = 1
    answer += next[0]
    for i in graph[next[1]]:
        if visit[i[1]] == -1:
            heapq.heappush(h, i)
print(answer)

 

위 문제는 알고리즘에서 매우 유명한 스패닝 트리에서의 최소 비용 구하기 문제이다.

 

스페닝 트리는 최소 경로 구하기와 달리, 모든 정점을 지나기만 하면 되기 때문에, 하나의 길로 연결될 필요가 없다는

특징을 가지고 있다.

 

또한 해당 간선의 가중치는 음수가 되도 된다는 조건이 있다.

 

그래프를 표현하기 위한 방법으로는 크게 2가지가 있다.

 

인접 행렬 표현 방식과, 인접 리스트 표현 방식이다.

 

인접 행렬 표현 방식은 필요없는 정보가 많이 들어가기 때문에, 정점이 많은 경우 불필요한 공간이

많다는 단점이 있다.

 

위 문제의 점점의 개수는 최대 10000개이다. 

 

따라서 나는 인접 리스트 표현 형식을 사용하여 그래프를 표현하였다.

 

graph = [[] for _ in range(v + 1)]
for _ in range(e):
    start, end, score = map(int, sys.stdin.readline().split())
    graph[start].append([score, end])
    graph[end].append([score, start])

 

 

 

 

최소 스패닝 트리를 풀기 위해서는 크게 크루스칼 알고리즘과 프림 알고리즘이 있다.

 

간단히 설명하면

 

프림 알고리즘은 시작 정점으로 부터 인접한 간선중 가장 작은 가중치를 가지는 간선을 선택하고, 

 

추가된 정점의 인접한 간선을 포함한 가능한 간선 중 가장 작은 간선을 선택하는 형식으로,

 

시작 정점으로부터, 인접한 정점을 하나씩 포함시켜가며, 포함된 정점과 이웃한 정점을 선택해,

점점 정점들을 늘려나가는 방식이다.

 

크루스칼 알고리즘은 시작 정점을 정하지 않고, 모든 간선 중 가장 작은 간선을 선택하고, 제외시킨다.

계속 가장 작은 간선을 선택하는데, 이때, 

1. 선택된 간선들 사이에서 사이클이 발생하면 안된다.

2. 이미 선택된 간선이어서는 안된다.

 

 

개인적으로 크루스칼 보다는 프림 알고리즘이 더 구현하기 쉽다고 생각해서 프림 알고리즘으로 구현했다.

 

프림알고리즘은 다음과 같은 순서로 구현했다.

1.모든 정점의 방문정보과, 총 점수, 시작 정보등을 초기화한다. 시작 정보 = [0,1] (가중치, 1번 노드부터 출발)

  (현재 1번 노드에서 출발하므로 가중치는 0이다.)

 

2. 시작 정보를 힙으로 만들어서 최소 가중치를 꺼낸다. (가중치를 0번째 인덱스에 넣은 이유)

  힙으로 pop된 노드는 다음 노드로 이동할 노드이다.

  

3. 선택된 다음 노드의 방문처리와, 총 스코어 계산 처리를 한다.

 

4. 다음 노드와 인접한 노드를 살피고, 방문하지 않았다면, 시작 정보(다음 돌아봐야 할 것)에 추가한다. heapq.heappush(...) 

 

1~4번을 반복하면 끝!

 

 

처음 while문을 들어가기 전에는 아무 것도 방문하지 않은 상태가 된다.

 

1.처음으로 while문에 들어가면, 1번 노드를 다음 노드로 선택한다.

2. 현재 1번 노드를 방문한 상태, 다음 인접한 노드를 살피면서 힙에 넣는다 (ex [2,3,4]

3. 두번째 while문에서 다음 노드를 선택한다. (ex 3번 노드가 가장 가중치가 적다고 가정)

4. 3번에서 뽑은 가장 작은 노드에 방문하고, 3번노드와 인접한 노드를 살피면서 힙에 넣는다. (ex[2,3,5,6]

.. 

힙에 아무것도 없을 때까지 반복한다.

 

처음에는 힙이 아니라 리스트에 집어넣으면서, 계산했지만, 시간 초과가 발생했다.

 

따라서 힙을 사용해서 가장 작은 가중치를 가지는 노드를 O(1) 만에 찾도록 하니 통과하였다.