[파이썬/python] 백준 - 1368 물대기

문제

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

---

문제 설명

1. N개의 논에 물을 대려고 한다. 물을 대는 방법에는 두 가지가 있다.  
   1. 직접 논에 우물을 판다.
   2. 이미 물을 대고 있는 다른 논으로부터 물을 끌어온다.
2. 각각의 논에 대해 우물을 파는 비용과, 논들 사이에 물을 끌어오는 비용들이 주어졌을 때, 모든 논에 물을 대는데 필요한 최소비용을 출력한다. 

 

---

풀이

0 2 2 2
2 0 3 3
2 3 0 4
2 3 4 0

위 같은 배열이 주어졌다고가정하자.

 

4개의 논에 물을 대야 한다.

예시 경우를 만들어 보면 , 1 -> 3 -> 2 -> 4, 1 -> 4 -> 2 -> 3 등 많은 경우의 수가 존재할 것 이다.

 

따라서 나는 A에서 B로 갈 수 있는 최소의 가중치를 찾는 것이라 생각했고, 플로이드 워셜로 초기 로직을 구현했다.

 

하지만 플로이드 워셜은 시작 정점과 도착 정점이 있을 때 최단거리를 구하는 것 이고

이 문제는 모든 정점을 최소의 가중치로 방문하는 문제였다. 즉 플로이드 워셜은 사용할 수 없다.

 

따라서 나는 최소의 가중치들로 모든 노드를 연결하는 알고리즘인 MST(최소 스패닝 트리)를 사용해서 문제를 풀이했다.

 

처음에 N개의 논이 최소의 가중치라면 여러 개를 선택할 수 있도록, 초기 N개의 논을 우물을 파는 가중치로 hq에 넣은 후

프림 알고리즘을 통해 해결했다.   

 

1. 각각의 논에 우물을 파는 가중치를 hq에 초기화한다.
2. 프림 알고리즘을 통해 최소 스패닝 트리를 구한다.
3. 최소 가중치를 출력한다. 

---

코드

import sys
import heapq
input = sys.stdin.readline
INF = sys.maxsize
N = int(input())

hq = [(int(input()),i) for i in range(N)]
cost = [list(map(int,input().split())) for _ in range(N)]
visited = [0 for _ in range(N+1)]
min_cost = 0
heapq.heapify(hq)
while hq:
    curr_cost,v1 = heapq.heappop(hq)
    if visited[v1]:
        continue
    min_cost += curr_cost
    visited[v1] = 1
    for v2 in range(N):
        if v1 == v2:
            continue
        heapq.heappush(hq,(cost[v1][v2], v2))
print(min_cost)

 

---

시간복잡도

• 프림 알고리즘은 hq를 사용하여 N개의 노드를 방문한다. 즉 각 노드를 방문했을 때 $ logV$ 만큼의 탐색을 진행하기 때문에 $O(ElogV)$의 시간복잡도가 소요된다.
• 최악의 경우 $N = 300$,$V = N*N = 9*10^4$이므로 약 $O(10^5)$의 시간복잡도가 소요된다.