백준 13418 [학교 탐방하기] 파이썬

https://www.acmicpc.net/problem/13418

 

# 백준 13418 학교 탐방하기
# 최소 스패닝 트리를 두 번 구해서 (최선/최악의 경우) 차이를 계산하는 문제
# - 좋은 길(C=0, 오르막길)은 피로도를 증가시킴
# - 나쁜 길(C=1, 내리막길)은 피로도를 증가시키지 않음
# 피로도 = (오르막길 개수)^2 이므로, MST를 두 번 구성해야 함

import sys
input = sys.stdin.readline

N, M = map(int, input().split())  # 건물 개수 N, 간선(길) 개수 M

# 유니온-파인드(Disjoint Set Union) 함수 정의
def find(x):
    """루트 노드를 찾는 함수 (경로 압축)"""
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    """두 노드를 합치는 함수 (다른 집합일 때만 합침)"""
    a = find(a)
    b = find(b)
    if a != b:
        parent[a] = b
        return True
    return False


# 간선 정보 입력 (M+1개, 0번째는 입구와 1번 건물 연결)
routes = []
for _ in range(M + 1):
    routes.append(list(map(int, input().split())))  # A, B, C(0=오르막, 1=내리막)


# -------------------------------
# 1) 최선(best case): 오르막길을 최소화하는 경우
#    => 내리막길(1)을 우선적으로 선택해야 하므로
#       내림차순 정렬 (C=1이 앞으로 오도록 정렬)
# -------------------------------
parent = [i for i in range(N + 1)]
best_ans = 0  # 선택된 오르막길 개수

routes.sort(key=lambda x: -x[2])  # C=1(내리막)이 앞으로 오게 정렬
for A, B, C in routes:
    if union(A, B):       # 사이클을 만들지 않는 경우만 선택
        if C == 0:        # 오르막길이면 카운트 증가
            best_ans += 1


# -------------------------------
# 2) 최악(worst case): 오르막길을 최대화하는 경우
#    => 오르막길(0)을 우선적으로 선택해야 하므로
#       오름차순 정렬 (C=0이 앞으로 오도록 정렬)
# -------------------------------
parent = [i for i in range(N + 1)]
worst_ans = 0

# 이미 routes를 내림차순 정렬했으므로
# 뒤에서부터 접근하면 오름차순 효과
for i in range(M, -1, -1):
    A, B, C = routes[i]
    if union(A, B):
        if C == 0:
            worst_ans += 1


# -------------------------------
# 최종 결과 출력
# 피로도 = (오르막길 개수)^2 이므로,
# 최악과 최선의 피로도 차이 출력
# -------------------------------
ans = (worst_ans + best_ans) * (worst_ans - best_ans)
print(ans)

최소 신장 트리(MST)가 익숙해진다.

크루스칼 알고리즘으로 구현하였다.