[BOJ] 13418. 학교 탐방하기

www.acmicpc.net/problem/13418

13418번: 학교 탐방하기

 

문제: 피로도가 최악인 경로와 최적인 경로의 피로도 차이를 구하라.

 

문제 이해하기!!

학교의 입구부터 학교 건물을 모두 소개해주는 문제입니다.

모든 경로를 탐방해야하므로 최소 스패닝 트리 문제라고 할 수 있습니다.

그러므로 이제 제일 잘 알려진 프림 알고리즘과 크루스칼 알고리즘 중 하나를 선택하여 풀면됩니다.

음... 혹시나 둘 다 모르는 분은 다익스트라가 익숙하시면 프림으로 하시고,

유니온파인드가 익숙하시면 크루스칼을 선택하시면 됩니다.

그마저도 모르겠다하면 프림을 먼저 공부해 보시는 걸 추천합니다^^

 

제가 푼 방식을 설명하자면,

1. 주어진 입력을 딕셔너리로 저장합니다. 다만, 0이 오르막길이므로 코스트를 XOR연산으로 넣어줍니다.

2. 힙큐로 최대 힙과 최소 힙을 사용하여 방문할 건물을 선택합니다.

3. 최대 힙으로 만든 길의 코스트의 제곱과 최소 힙으로 만든 길의 코스트의 제곱의 차를 출력합니다.

 

여기서 주의할 점이 두 가지 있습니다.

1. 1번과 2번이 연결되어 있을 때 사이가 오르막길이면 1번에서 2번으로 갈 때도 오르막길로 치고, 2번에서 1번으로 갈 때도 오르막길로 치지 않아요. 즉, 일반적인 상식처럼 한쪽이 오르막길이라고 반대쪽을 내리막길로 생각하면 안 됩니다.

2. 이미 방문한 것을 제외시키지 않는다면 시간초과가 나게 됩니다. - 여기선 코스트가 1과 0 밖에 없기 때문에 최대 1000 * 1000 * 500의 시간이 나올 수 있습니다. 조심하며 시간초과를 어떻게 없앨지 생각해보세요!

1번은 아예 문제자체를 잘 못 생각할 수도 있으니 정말! 주의 바랍니다!

 

코드:

from collections import defaultdict
from sys import stdin
input = stdin.readline
import heapq

N, M = map(int, input().split())
dic = defaultdict(list)

# 0이 오르막길이므로 코스트를 XOR연산으로 변환
for i in range(M + 1):
    a, b, c = map(int, input().split())
    dic[a].append((c^1, b))
    dic[b].append((c^1, a))

# MST 함수 (프림)
def MST(Q, ch):
    res = 0
    heapq.heapify(Q)
    visit = [False] * (N + 1)
    visit[0] = True
    while Q:
        cost, next = heapq.heappop(Q)
        # 이미 방문한 것은 제외
        if visit[next]:
            continue
        res += cost
        visit[next] = True
        for i in dic[next]:
            # 이미 방문한 것은 제외(여기서 안해주면 시간초과)
            if visit[i[1]]:
                continue
            heapq.heappush(Q, (i[0] * ch, i[1]))
    return res

# 처음 코스트의 부호변환 작업(최대힙을 위해)
adj = [(-c, n) for c, n in dic[0]]

# 최대와 최소 값의 차 구함.
print(MST(adj[:], -1)**2 - MST(dic[0][:], 1)**2)