문제
N개의 논에 물을 대려고 한다. 각 논에 직접 우물을 파거나 다른 논에서 수로를 연결할 수 있을 때, 모든 논에 물을 대는 최소 비용을 구하라.
입력
첫째 줄에 논의 수 N, 이후 N줄에 각 논의 우물 비용, 이후 N×N 행렬로 논 사이 수로 비용이 주어진다.
출력
최소 비용을 출력한다.
예제
| 입력 | 출력 |
|---|---|
3 5 4 4 0 2 2 2 0 3 2 3 0 | 9 |
풀이
가상의 0번 노드를 추가하여 우물 비용을 0번 노드와의 간선으로 표현하면, MST 문제로 변환할 수 있다.
- 각 논의 우물 비용을 가상 노드 0과 해당 논 사이의 간선으로 추가한다
- 논 사이 수로 비용을 간선으로 추가한다 (중복 방지를 위해
i < j인 쌍만) - 모든 간선을 비용 오름차순으로 정렬한다
- 크루스칼 알고리즘으로 N개의 간선을 선택하여 MST를 구성한다
- Union-Find로 사이클을 검사하며 간선을 추가한다
핵심 아이디어: 우물을 파는 행위를 "가상 수원 노드에 연결"로 모델링하면, 우물과 수로를 통합한 하나의 MST 문제로 귀결된다.
코드
import sys
input = sys.stdin.readline
MAX = float("inf")
N = int(input())
road_list = list()
for i in range(N):
well_cost = int(input())
road_list.append((well_cost, i + 1, 0))
for i in range(N):
road_cost = list(map(int, input().split()))
for j in range(i + 1, N):
road_list.append((road_cost[j], i + 1, j + 1))
road_list.sort()
parent = list(range(N + 1))
def find(a):
if a == parent[a]:
return parent[a]
parent[a] = find(parent[a])
return parent[a]
def union(pa, pb):
if pa < pb:
parent[pb] = pa
else:
parent[pa] = pb
answer, cnt = 0, N
for cost, a, b in road_list:
if not cnt:
break
pa = find(a)
pb = find(b)
if pa != pb:
cnt -= 1
union(pa, pb)
answer += cost
print(answer)복잡도
- 시간: O(E * log E) (E: 간선 수, 최대 N^2)
- 공간: O(N^2)