[ 백준 ][ 골드2 ] 1368번 - 물대기 ( 파이썬 PYTHON3 )

Heeto 2023. 3. 11. 16:42

링크 : https://www.acmicpc.net/source/57235108

로그인

www.acmicpc.net

코드

'''
먼저 처음 물을 대려면 우물을 먼저 하나 파야한다.

이 우물을 어디에 먼저 팔 것인가는 최소 스패닝 트리 문제 특성상 그리디하게 선정할 수 있을 것 같은데
어떤 기준으로 그리디하게 선정할 것인가가 문제이다.

1. 다른 논에서 끌어올 때 가장 많은 비용이 드는 것.
2. 우물을 팔 때 가장 적은 비용이 드는 것.

1번으로 했을 떄는 우물을 팔 떄 비용이 훨씬 더 엄청나게 많이 들 수도 있다는 단점이 존재한다.
2번으로 했을 때는 해당 우물은 적은 비용으로 팔 수 있을 지 몰라도 다른 논에 물을 댈 때 엄청나게 많은 비용이 발생할 수 있다.

그렇다면 2번으로 먼저 한 후에 어떤 논에 물을 댈 때 이미 만들어진 논들에서 끌어오는 가장 적은 비용과 우물을 새로 파는 것 중 선택한다면 어떨까?

근데 생각해보니 논에 물을 대는 순서가 상관이 있는 것일까?
어차피 모든 논에 물이 연결될 것이라면 어떤 것을 먼저 하는지는 상관 없이 그냥 우물을 파는 것과 다른 논에서 끌어오는 가장 적은 비용 중 작은 것을 선택하면 되는 게 아닐까?
만약 A논에 B의 물을 끌어오고 싶은데 현재 순서상 B는 아직 물을 공급하지 않았을 때, 어떤 방식으로든 순서가 바뀌면 B가 먼저 공급될 수 있는 게 아닐까?

여기서 또 생각해보면 단순히 적은 비용을 선택하다보니 A는 B에서 물을 끌어오고 B도 A에서 물을 끌어오게 될 수도 있다. 이런 경우는 싸이클이므로 싸이클은 배제하는 코드가 필요하다.
싸이클을 배제하려면 비용이 정렬된 상태에서 적은 비용들 먼저 처리를 하면서 싸이클이 발생하지 않는 선택을 해야 한다.

그렇다면 정렬은 어떤 것을 기준으로 해야 할까?
정렬은 역시 비용을 기준으로 한 번에 처리해야 한다.
한 번에 처리하기 위해서는 우물을 파는 방법을 0번 노드에서 다른 노드로의 간선으로 여기고 최소힙을 사용하면 될 것이다.
'''

import heapq

n = int(input())
heap = []
parent = [i for i in range(n+1)]

for i in range(1,n+1):
    heapq.heappush(heap, [int(input()),0,i])

for i in range(1,n+1):
    cost = list(map(int,input().split()))
    for j in range(1,n+1):
        if i != j:
            heapq.heappush(heap, [cost[j-1], i, j])

def find(node):
    if parent[node] != node:
        parent[node] = find(parent[node])
    return parent[node]

def union(a,b):
    parent_a, parent_b = find(a), find(b)
    if parent_a != parent_b:
        parent[parent_b] = parent_a
        return True
    return False

ans = 0
while heap:
    cost, a, b = heapq.heappop(heap)
    if union(a,b):
        ans += cost

print(ans)

풀이

최소 신장 트리(MST) 문제이고 크루스칼 알고리즘을 사용해서 해결했다.

주석은 대충 내가 문제를 읽으면서 생각한 내용이다.

'크루스칼 알고리즘'은 그래프 내의 모든 정점들을 가장 적은 비용으로 연결하기 위한 알고리즘이다.

간선들을 비용에 따라 오름차순으로 정렬한다.
간선이 잇는 두 정점의 root노드를 찾는다.
두 정점의 root가 다르다면 바꾸어서 연결시켜주고 비용을 더한다.

이와 비슷하게 '프림 알고리즘'이 있다.

임의의 정점을 선택한다.
해당 정점에서 갈 수 있는 간선을 최소힙(비용기준)에 넣는다.
최소힙에서 값을 꺼내어 해당 정점을 아직 방문하지 않았다면 방문처리하고 비용을 넣는다.

그래서 이 문제는 다른 문제들처럼 간선이 주어지지는 않지만 모든 논들끼지 이어져있고 모든 논들끼리 비용이 있기 때문에 간선이 많다고 볼 수 있다.

하지만 그 간선들을 최소힙에 넣어서 정렬을 시키지 않고 최소값들만 뽑아서 사용한다면 시간복잡도가 크게 증가하지 않을 것이다.

사실 프림 알고리즘을 어떻게 적용시키는지 잘 몰라서 크루스칼로 했다.

먼저 힙(우선순위큐)와 부모배열을 초기화한다.

부모배열은 처음엔 자기 자신을 부모로 삼는다.

N번 우물을 직접 팔 때의 비용이 주어지는데 이를 0->i의 간선으로 여기고 최소힙에 저장한다.

또 N번 i번째 논, j번째 논에 대하여 I->J == J->I의 비용이 주어지므로 최소힙에 모든 논과 논사이의 거리와 두 논의 정보를 저장한다.

이렇게 하면 우물을 파던 다른 논에서 물을 길어오던 상관없이 비용에 따라 정렬된 최소값들을 뽑아낼 수 있다.

이때부터는 다른 MST문제와 다를 것 없이 두 노드의 부모가 다르면 합쳐주고 비용을 누적시켜주면 된다.

다른 사람들 코드를 보니 프림 알고리즘을 사용하면 70ms 이내로 엄청나게 빠른 알고리즘을 짤 수 있는 것 같다.

참고로 위의 크루스칼 알고리즘은 404ms가 나왔다.

프림 알고리즘으로 짠 코드를 한 번 봐보자.

import sys, heapq
input = sys.stdin.readline

n = int(input())
w, q = [], []
for i in range(n):
    cost = int(input())
    w.append(cost)
    heapq.heappush(q, (cost, i))
    
p = [list(map(int, input().split())) for _ in range(n)]
cost = 0
visited = [0] * n

while q:
    x, cur = heapq.heappop(q)
    if visited[cur] == 0:
        visited[cur] = 1
        cost += x

        for nxt in range(n):
            if cur != nxt:
                if w[nxt] > p[cur][nxt]:
                    w[nxt] = p[cur][nxt]
                    heapq.heappush(q, (w[nxt], nxt))
print(cost)

w배열을 만들어서 우물일 파는 방법을 따로 저장하고 최소힙에 넣어주었다.

또 이차원배열로 주어지는 물을 길어오는 비용들은 그대로 저장하고 방문처리배열을 생성했다.

우물을 파는 방법에서 가장 적은 비용이 드는 것부터 빼서 해당 노드가 방문처리되지 않았다면 방문처리하고 비용을 더해준다.

또 BFS처럼 현재 방문처리한 노드에서 갈 수 있는 노드들 중 우물을 파는 것보다 물을 길어다 주는 것이 더 비용이 적다면 최소힙에 넣어준다.

프림알고리즘은 BFS와 꽤 유사한 부분들이 많고 우물을 파는 것과 물을 길어오는 것의 통합을 while문에서 하게 된다.

어찌보면 내가 주석으로 적은 내용은 이 프림 알고리즘에 더 가까웠다는 생각이 든다.

시간복잡도에서 많은 차이를 보이니 역시 MST에 대해서도 한 가지보다는 여러가지 방법을 모두 익혀야 겠다.