[백준/Python]1368: 물대기

문제

https://www.acmicpc.net/problem/1368
백준 문제집 - 0x1B강 - 최소 신장 트리

 

알고리즘: 최소 스패닝 트리

---

풀이

 

어떤 알고리즘을 사용하느냐에 따라 풀이 방식이 달라진다.

 

1️⃣ 크루스칼 알고리즘 사용

크루스칼 알고리즘은 기본적으로 (노드1, 노드2, 비용)형태의 데이터를 오름차순으로 정렬하여 사용한다.
하지만 이 문제에서는 두 가지 조건이 제시된다.

1) 다른 논에서 물을 끌어올것인지

2) 논에 새로운 우물을 팔 것인지

→ 일단 정렬 후 연결비용과 우물비용 중에서 비교해야하나? 싶었지만 아니었다.

 

🗝️ 가상의 논 N을 설정한 뒤, (i, N, 논 i의 우물비용)을 그래프에 추가로 저장해주면 된다.
그럼 모든 논과 가상의 논 사이에 우물을 파는 비용으로 간선이 연결되고, 크루스칼 알고리즘을 적용할 수 있다.

• 만약 논 i에 직접 우물을 파는것이 더 저렴하다면, (i, N)간선이 선택됨
• 다른 논 j로부터 물을 끌어오는것이 더 저렴하다면, (i, j)간선이 선택됨

이런식으로 자연스럽게 1) 물을 끌어오는비용, 2) 우물을 파는 비용 중 더 작은 값을 선택할 수 있다.

N에 대한 추가 간선 이외에는 기존 크루스칼 로직과 동일하다.

 

전체 코드

# 1. 크루스칼 알고리즘 풀이
# 메모리: 38240KB / 시간: 88ms
from sys import stdin


input = stdin.readline

N = int(input())
water = [int(input()) for _ in range(N)]

# 가상의 논 N을 설정해야할듯.
# (i, N) => i번째 논에서 직접 우물을 파는 경우
graph = [(i, N, water[i]) for i in range(N)]

for i in range(N):
    line = list(map(int, input().split()))
    for j in range(i):
        graph.append((i, j, line[j]))

parent = list(range(N+1))

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


def union(a, b):
    a, b = find(a), find(b)

    if a != b:
        if a < b:
            parent[b] = a
        else:
            parent[a] = b
        return True
    return False


ret = 0
graph.sort(key=lambda x: x[2])

for u, v, w in graph:
    if union(u, v):
        ret += w


print(ret)

 

2️⃣ 프림 알고리즘 풀이

프림 알고리즘은 시작정점부터 단계적으로 트리를 확장해 나가는 방식이다.
따라서 프림 알고리즘을 사용할 경우 우물비용이 가장 작은 논을 먼저 파둬야한다.
그리고 해당 논을 기준으로 MST를 구성하는데, (현재까지의 최소비용, 우물비용, 끌어오는 비용) 중 최소값을 선택하면 된다.

(리스트 or heapq 모듈을 사용해서 풀 수 있다.)

 

리스트 사용

과정은 다음과 같다.

1. 가장 적은 우물 비용을 가진 논을 선택
2. 해당 논의 우물 비용을 초기 비용으로 설정
3. 전체 N개의 논에 대해 반복
   • 아직 방문하지 않은 논 중 최소 비용을 가진 논을 선택
   • 해당 논을 방문처리 후 전체 비용에 더해줌
   • 선택한 논에서 연결할 수 있는 다른 논들의 비용을 업데이트
     • min(현재까지의 해당 논 최소비용, 우물비용, 연결비용)

 

 

전체 코드

# 2. 프림 알고리즘 풀이 (리스트)
# 메모리: 35480KB / 시간: 72ms
from sys import stdin


input = stdin.readline

N = int(input())
water = [int(input()) for _ in range(N)]
graph = [tuple(map(int, input().split())) for _ in range(N)]
cost = [float("inf")] * N
visited = [False] * N

# 시작점 찾기 (가장 작은 우물 비용)
start = -1
start_water = float("inf")

for i in range(N):
    if water[i] < start_water:
        start = i
        start_water = water[i]

cost[start] = start_water
ret = 0

for _ in range(N):
    min_cost = float("inf")
    min_node = -1

    for i in range(N):
        if not visited[i] and cost[i] < min_cost:
            min_cost = cost[i]
            min_node = i

    visited[min_node] = True
    ret += min_cost

    for nxt in range(N):
        if visited[nxt]:
            continue
        # 각각의 비용을 독립적으로 비교
        cost[nxt] = min(cost[nxt], graph[min_node][nxt], water[nxt])

print(ret)

 

우선순위 큐(heapq) 사용

리스트 구현 방식과 동일하지만, 최소 비용을 찾는 과정을 힙으로 최적화한것이다.
과정 역시 거의 동일함.

1. 가장 적은 우물 비용을 가진 논을 선택
2. heap에 (비용, 논 번호) 삽입
3. 비용이 적은 경우부터 하나씩 꺼내며 체크함
   • 이미 방문한 노드라면 건너뜀
   • 아니라면 방문처리 후 힙에서 꺼낸 비용을 전체 비용에 추가
   • 노드 수 카운트 +1
     • 만약 카운트한 노드 수가 N 이상이라면 최소 신장 트리를 완성한 셈이므로 break
   • (다른 논과 연결하는 비용, 우물을 파는 비용) 중 더 작은 값을 선택
   • 해당 값이 현재까지 해당 논의 최소비용보다 작다면 업데이트 후 힙에 추가

 

전체 코드

# 3. 프림 알고리즘 풀이 (heapq)
# 메모리: 37556KB / 시간: 72ms
from sys import stdin
from heapq import heappush, heappop


input = stdin.readline

N = int(input())
water = [int(input()) for _ in range(N)]
graph = [tuple(map(int, input().split())) for _ in range(N)]

costs = [float("inf")] * N
visited = [False] * N

# 시작점 찾기 (가장 작은 우물 비용)
start = -1
start_water = float("inf")

for i in range(N):
    if water[i] < start_water:
        start = i
        start_water = water[i]

heap = [(start_water, start)]

costs[start] = start_water
ret = cnt = 0

while heap:
    cost, node = heappop(heap)

    if visited[node]:
        continue

    visited[node] = True
    ret += cost
    cnt += 1

    if cnt >= N:
        break

    for nxt, cost in enumerate(graph[node]):
        if visited[nxt]:
            continue

        # 다른 논과 연결하는 비용, 우물을 파는 비용 중 더 작은 값을 선택
        new_cost = min(cost, water[nxt])

        # 해당 값이 최소비용보다 작다면 업데이트 후 힙에 추가한다.
        if new_cost < costs[nxt]:
            heappush(heap, (new_cost, nxt))
            costs[nxt] = new_cost


print(ret)

---

개인적으론 크루스칼을 사용한 풀이가 마음에 든다.

N번째 논으로 설정하니 참 간단하다...