[백준 13418] 학교 탐방하기

문제 출처 : https://www.acmicpc.net/problem/13418

13418번: 학교 탐방하기

이 문제는 최소 신장 트리를 두 번 구성해야 하는 문제 입니다. 최소 신장 트리에 대해 잘 모른다면 아래 링크를 통해 최소 신장 트리에 대해 이해하시기 바랍니다.

https://wikidocs.net/207011

05. 최소 신장 트리

먼저 오르막길을 가장 많이 갈 수 있는 방법으로 한 번 구하고, 다음으로 내리막 길을 가장 많이 갈 수 있는 방법으로 구해 두 값의 제곱의 차로 답을 구할 수 있습니다. 우리는 크루스칼 알고리즘을 통해 이 문제를 해결해 보겠습니다.

내리막 길은 1, 오르막 길은 0으로 표시되어 있기 때문에 이 가중치를 통해 최악의 피로도, 최선의 피로도는 쉽게 구할 수 있습니다.

문제 풀이

그럼 문제를 직접 풀어보도록 하겠습니다.

입력 받기

mii = lambda : map(int, input().split())
N, M = mii()

arr = []
for _ in range(M + 1):
    a, b, c = mii()
    arr.append((c, a, b))

arr.sort()

건물의 수 N과 도로의 개수 M을 입력 받습니다. 다음으로 모든 도로의 정보를 입력 받았습니다. 이 때 도로가 오르막 길인지 내리막 길인지 알 수 있는 가중치 c의 값을 첫 번째로 받아 이 값으로 정렬이 될 수 있게 하였습니다. 그리고 입력받은 도로 arr을 정렬해 주었습니다.

유니온 파인드

크루스칼 알고리즘은 유니온 파인드를 통해 노드들이 사이클을 가지고 있는지 없는지를 판단합니다. 최소 신장 트리는 트리로 구성되기 때문에 사이클이 있으면 안됩니다.

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

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

    if a < b:
        parent[b] = a
    else:
        parent[a] = b

유니온 파인드는 정형화되어 있기 때문에 따로 설명하지 않겠습니다. 잘 이해가 안된다면 아래 링크를 통해 유니언 파인드에 대해 다시 한번 학습해 보시기 바랍니다.

https://wikidocs.net/206632

02. 유니온 파인드

최악의 피로도 계산하기

max_ans = 0
parent = list(range(N + 1))
for c, a, b in arr:
    if find(a) == find(b):
        continue

    union(a, b)
    max_ans += 1 if c == 0 else 0

max_ans **= 2

먼저 최악의 피로도를 계산하겠습니다. 앞서 정렬된 arr은 오름차순 정렬이 되어 있습니다. 오르막 길이 0, 내리막 길이 1이기 때문에 오르막 길을 우선적으로 연결합니다. 따라서 이 값이 최악의 피로도를 계산할 수 있게 됩니다. 유니온 파인드로 연결 정보를 확인하여 피로도를 계산한 것입니다. 오르막 길의 피로도를 모두 더한 다음에 마지막에 그값을 제곱하여 값을 구합니다.

최선의 피로도 계산하기

min_ans = 0
parent = list(range(N + 1))
for c, a, b in arr[::-1]:
    if find(a) == find(b):
        continue

    union(a, b)
    min_ans += 1 if c == 0 else 0

min_ans **= 2 

print(max_ans - min_ans)

다음으로 최선의 피로도를 구해보겠습니다. 유니온 파인드를 다시 해야 하기 때문에 parent를 다시 초기화 하였습니다. 그리고 정렬된 도로 arr을 거꾸로 하였습니다. arr[::-1]에 대해 잘 모르면 아래 링크를 통해 리스트를 거꾸로 하는 방법에 대해 확인해 보시기 바랍니다.

https://wikidocs.net/192334#2-

01. 리스트 기본 사용방법

마지막으로 최악의 피로도에서 최선의 피로도를 빼서 출력하면 됩니다.

전체 코드

전체 코드를 확인해 보겠습니다.

mii = lambda : map(int, input().split())
N, M = mii()

arr = []
for _ in range(M + 1):
    a, b, c = mii()
    arr.append((c, a, b))

arr.sort()
def find(a):
    if a != parent[a]:        
        parent[a] = find(parent[a])
    return parent[a]

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

    if a < b:
        parent[b] = a
    else:
        parent[a] = b

max_ans = 0
parent = list(range(N + 1))
for c, a, b in arr:
    if find(a) == find(b):
        continue

    union(a, b)
    max_ans += 1 if c == 0 else 0

max_ans **= 2

min_ans = 0
parent = list(range(N + 1))
for c, a, b in arr[::-1]:
    if find(a) == find(b):
        continue

    union(a, b)
    min_ans += 1 if c == 0 else 0

min_ans **= 2
print(max_ans - min_ans)