[Python/파이썬] 백준 1368번 물대기

1368번: 물대기

 

문제

선주는 자신이 운영하는 N개의 논에 물을 대려고 한다. 물을 대는 방법은 두 가지가 있는데 하나는 직접 논에 우물을 파는 것이고 다른 하나는 이미 물을 대고 있는 다른 논으로부터 물을 끌어오는 법이다.

각각의 논에 대해 우물을 파는 비용과 논들 사이에 물을 끌어오는 비용들이 주어졌을 때 최소의 비용으로 모든 논에 물을 대는 것이 문제이다.

입력

첫 줄에는 논의 수 N(1 ≤ N ≤ 300)이 주어진다. 다음 N개의 줄에는 i번째 논에 우물을 팔 때 드는 비용 Wi(1 ≤ Wi ≤ 100,000)가 순서대로 들어온다. 다음 N개의 줄에 대해서는 각 줄에 N개의 수가 들어오는데 이는 i번째 논과 j번째 논을 연결하는데 드는 비용 Pi,j(1 ≤ Pi,j ≤ 100,000, Pi,j = Pj,i, Pi,i = 0)를 의미한다.

출력

첫 줄에 모든 논에 물을 대는데 필요한 최소비용을 출력한다.

코드

from heapq import heappop, heappush

n = int(input())

parent = [i for i in range(n+1)]
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    a = find(a)
    b = find(b)
    if a == b:  # 사이클을 형성
        return
    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        
costs = []  # i번째 논에 우물 파는데 드는 비용 = 0번(가상의 논)에서 i번으로 물 대주는 비용
for i in range(1, n+1):
    heappush(costs, (int(input()), 0, i))
    
# i번째 논과 j번째 논을 연결하는데 드는 비용
for i in range(1, n+1):  
    arr = [0] + list(map(int, input().split()))
    for j in range(i+1, n+1):
        heappush(costs, (arr[j], i, j))
ans = 0
cnt = 0
while costs:
    c, a, b = heappop(costs)
    # 연결하는데 드는 비용
    if find(a) != find(b):
        union(a,b)
        ans += c
        cnt += 1
        if cnt == n:
            break        
print(ans)

직접 물을 퍼올리는 파는 논도 결국 어디선가 물을 끌어온다고 생각할 수 있으므로 논에 우물을 파는 경우는 가상으로 존재하는 0번 논에서 끌어온다고 생각하면 된다. 그러므로 i번째 논에 우물을 팔 때 드는 비용은 0번 논에서 i번째 논을 연결하는 비용이라고 바꿔서 생각하자.

(i번째와 j번째 논 연결비용, i, j) 형태로 최소힙에 넣어주고 크루스칼 알고리즘을 이용하여 MST를 만들어서 모든 논에 물을 대는 가장 최소의 비용을 구하였다.

 

처음에는 가상의 노드를 설정할 생각을 하지 못하고 아래와 같이 풀었었다.

from heapq import heappop, heappush

n = int(input())

parent = [i for i in range(n+1)]
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    a = find(a)
    b = find(b)
    if a == b:  # 사이클을 형성
        return
    if dig_costs[a] <= dig_costs[b]:    # 우물 파는데 드는 비용이 더 적은 쪽이 부모
        parent[b] = a
    else:
        parent[a] = b
        
dig_costs = []  # i번째 논에 우물 파는데 드는 비용
for _ in range(n):
    dig_costs.append(int(input()))
connect_costs = []
hq = []
# i번째 논과 j번째 논을 연결하는데 드는 비용
for i in range(n):  
    arr = list(map(int, input().split()))
    connect_costs.append(arr)
    for j in range(i+1, n):
        heappush(hq, (arr[j], i, j))
ans = sum(dig_costs)
cnt = 0
while hq:
    c, a, b = heappop(hq)
    root_a = find(a)
    root_b = find(b)
    # 연결하는데 드는 비용
    if root_a != root_b:
        # a, b를 연결하는게 더 유리한 경우
        connect_cost = connect_costs[a][b] - max(dig_costs[root_a], dig_costs[root_b])
        if connect_cost < 0:
            union(a,b)
            ans += connect_cost
        cnt += 1
    if cnt == n-1:
        break        
print(ans)

크루스칼 알고리즘을 이용한 풀이.

a, b를 연결하는게 유리한 경우는 연결하는 비용이 음수가 되는 경우라고 생각했다. 연결하지 않는 경우에는 비용에 변화가 없다 생각하여 아무것도 해주지 않고 cnt만 증가시켜주었다.

union을 할 때는 그 그래프에서 우물을 파는데 가장 최소의 비용을 갖는 논이 루트 노드가 될 수 있도록 해주었다.

 

그런데 이렇게 풀으니 그때는 연결 안 하는게 유리했는데 나중에 보면 연결하는게 더 유리한 경우가 발생할 수 있을 것 같다. 그래서 테스트케이스는 돌아갔지만 결과는 틀렸다고 뜬 거 같다.