본 포스팅은 해당 문제를 프림 알고리즘으로 푼 풀이에 대해 정리하였다.
그러나 크루스칼 알고리즘으로도 풀 수 있기 때문에 아래 포스팅도 참고 바란다.
더해서 아래 포스팅에 더 쉽게 이해할 수 있는 설명이 되어 있는 것 같아 본 포스팅을 이해했어도 아래 포스팅도 참고하는 것을 추천한다.
📋 참고 포스팅:[백준/Python] 1368. 물대기 - 크루스칼
백준
난이도 : Gold 2
문제 제목 : 물대기
import sys
import heapq
# 최소 신장 그래프
input = sys.stdin.readline
n = int(input())
self_cost = [int(input()) for _ in range(n)] # 직접 우물을 파는 비용
cost = [list(map(int, input().split())) for _ in range(n)] # 다른 논으로부터 물을 끌어오는 비용
tmp = [int(1e9), 0] # 처음 팔 논의 (비용, 정점)
# 직접 우물 파는 비용 중 최소 비용과 논 찾기
for i in range(n):
if self_cost[i] < tmp[0]:
tmp[0] = self_cost[i]
tmp[1] = i
heap = [tmp]
visited = [0] * n # 해당 논에 물을 댔는지 관리
cnt = 0 # 물을 댄 논 개수
result = 0 # 최종 비용
while heap:
if cnt == n:
break
c, v = heapq.heappop(heap) # 힙에 있는 최소 비용, 논 pop
if visited[v]: # 이미 물을 댄 논이면 pass
continue
visited[v] = 1
cnt += 1
result += min(self_cost[v], c) # 직접 해당 논에 우물을 파는 것과 c의 비용 비교
for i in range(n): # 이번 턴에 물을 댄 논과 연결된 논들을 heap에 추가
heapq.heappush(heap, [cost[v][i], i])
print(result)
최소 신장 트리를 구하는 프림 알고리즘을 살짝 변경한 풀이이다.
위 풀이는 오답이므로, 아래 '✏️ 풀이 2'의 풀이부터 참고 바란다. 그 후 더 공부해보고 싶다면 '✏️ 풀이 2'의 풀이와 비교해서 위 풀이의 문제점을 찾아봐도 좋을 것 같다.
간단히 코드를 설명해보겠다. 전체적으로 '✏️ 풀이 2'와 비슷하다.
직접 논의 위치에 우물을 팔 수도 있어서 모든 정점이 연결되지 않을 수 있음을 주의해야 한다.
우선 처음 직접 해당 위치에 우물을 팔 시작점을 구해야 하는데, 이는 직접 우물을 파는 비용이 최소인 논으로 결정한다. (어차피 결과적으로 모든 논을 거쳐야하고, 최소 1개의 논은 직접 우물을 파야하기 때문에 직접 우물을 파는 비용이 최소인 논을 시작점으로 두었다.)
그리고 heap에 시작점을 넣고, 비용이 최소인 것을 빼면서 프림 알고리즘을 수행해 나간다.
적절한 heap의 요소((비용, 정점))를 pop하면, 해당 논에 우물을 직접 파는 것과 pop한 비용을 비교하여 더 작은 값을 최종 비용에 추가한다.
-> 이 부분이 일반적인 프림 알고리즘과 다른 점으로, result += min(self_cost[v], c) 부분이다.
✅ 오답 이유:
이 풀이는 일부 테스트 케이스에서 오답이 출력된다.
'✏️ 풀이 2'와 비교하여 다른 점을 살펴보면 해당 테스트 케이스를 발견할 수 있다.
🍎 반례

위와 같은 케이스에서 정답은 다음과 같을 것이다.

그러나 이 풀이에서는 다음과 같은 답이 나올 것이다.

이러한 답이 나온 이유는, 직접 논에 우물을 파는 2에 대한 요소가 heap에 추가되기 전에,
비용이 더 큰 다른 논과의 연결로 물이 대지기 때문이다.
위의 예시에서 heap에서 요소가 pop되는 올바른 순서와 이 풀이의 순서를 살펴보면 다음과 같다. (위에서부터 반시계 방향으로 인덱스를 부여) 직접 그려보면서 확인해보는 것을 추천한다.
(1, 0), (10, 1), (2, 2)(1, 0), (10, 1), (2, 2)pop되는 요소: (1, 0)(4, 1), (5, 2)(10, 1), (2, 2), (4, 1), (5, 2)pop되는 요소: (2, 2)(5, 0), (1, 1)(10, 1), (4, 1), (5, 2), (5, 0), (1, 1)pop되는 요소: (1, 1)(4, 0), (1, 1)(10, 1), (4, 1), (5, 2), (5, 0), (4, 0), (1, 1)(1, 0)(1, 0)pop되는 요소: (1, 0)(4, 1), (5, 2)(4, 1), (5, 2)pop되는 요소: (4, 1)(4, 0), (1, 1)(5, 2), (4, 0), (1, 1)pop되는 요소: (1, 1)(5, 0), (1, 1)(5, 2), (4, 0), (5, 0), (1, 1)✅ 이러한 문제를 해결하기 위해서는, 초반에 모든 논에 대해 직접 논에 우물을 파는 비용과 인덱스를 heap에 추가해야 한다.
import sys
import heapq
# 최소 신장 그래프
input = sys.stdin.readline
n = int(input())
self_cost = [int(input()) for _ in range(n)] # 직접 우물을 파는 비용
cost = [list(map(int, input().split())) for _ in range(n)] # 다른 논으로부터 물을 끌어오는 비용
heap = [(self_cost[i], i) for i in range(n)] # 직접 우물을 파는 경우에 대해 (비용, 정점) 순으로 모두 추가한다.
heapq.heapify(heap) # list인 heap을 heapq로 변경한다.
visited = [0] * n # 해당 논에 물을 댔는지 관리
cnt = 0 # 물을 댄 논 개수
result = 0 # 최종 비용
while heap:
if cnt == n:
break
c, v = heapq.heappop(heap) # 힙에 있는 최소 비용, 논 pop
if visited[v]: # 이미 물을 댄 논이면 pass
continue
visited[v] = 1
cnt += 1
result += c
for i in range(n): # 이번 턴에 물을 댄 논과 연결된 논들을 heap에 추가
heapq.heappush(heap, (cost[v][i], i))
print(result)
✅ 풀이 한줄 설명:
최소 신장 트리를 구하는 프림 알고리즘을 살짝 변경한 풀이이다.
직접 논의 위치에 우물을 팔 수도 있어서 결과적으로 모든 정점이 연결되지 않을 수 있음을 주의해야 한다.
만약 최소 신장 트리 알고리즘을 잘 모른다면 먼저 공부해보기를 추천한다. (위 풀이에서는 크루스칼 알고리즘과 프림 알고리즘 중 프림 알고리즘을 사용하였다. 그러나 둘 다 공부해오기를 추천한다.)
✅ 풀이 자세한 설명:
heap 초기화하기🍎 1. heap 초기화하기
우선 처음에는 최소 하나의 논에 직접 우물을 파야 알고리즘을 시작할 수 있다. 따라서 모든 논에 대해, 직접 논에 우물을 파는 비용과 해당 논의 인덱스를 (비용, 정점) 형식으로 heap에 추가한다.
이렇게 하면 최소 비용을 가지는 요소부터 pop되어 2번 과정이 시작된다.
🍎 2. 프림 알고리즘 수행하기
최소 신장 트리의 프림 알고리즘을 수행한다.
heap에 있는 요소 중 비용이 가장 작은 요소를 pop하면서 반복문을 수행한다.
이번 턴의 pop한 논에 이미 물을 댔다면 continue를 하고, 그렇지 않다면 물을 대준다.(visited, cnt, result 등 업데이트)
그리고 물을 대준 경우, 해당 논과 연결된 다른 논들을 (비용, 정점) 형식으로 heap에 추가한다.
해당 문제, 풀이에 대한 GitHub Repository 링크는 다음과 같다.
GitHub - 백준(Gold) '1368. 물대기'
GitHub - [27강] 최소 신장 트리/연습문제 '1368. 물대기'
해당 문제를 풀 때 프림 알고리즘을 어떻게 변형해야 할지 고민을 많이 했었는데,
최종적으로 답은 최소 신장 트리 구하는 일반적인 프림 알고리즘과 별 차이가 없었다.
그래도 여러 고민을 해보면서 최소 신장 트리 알고리즘을 더 깊게 공부할 수 있었던 것 같다.