백준 13418 <학교 탐방하기> Python

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

---

최소 스패닝 트리 문제입니다.

다만 그닥 추천드리고 싶지 않은 문제입니다.

 

문제 자체는 그냥 크루스칼 알고리즘을 두 번 사용해 풀면 되기에 어렵지는 않지만 입력과 조건이 이상하다고 생각됩니다.

분명 도로의 개수 M 이라고 작성되어 있지만 당장 입력에는 M+1 개의 도로가 주어집니다.

다 풀어놓고도 입력 오류를 이유로 틀린 사람이 많을 것 같습니다.

 

코스트와 피로도가 반전되어 있다는 점과 입력이 이상하다는 점만 유의하면 쉽게 풀 수 있습니다.

---

import sys

input = sys.stdin.readline


## Union-Find

# Find
def Find(group_list, node):
    # 루트 노드가 아니면
    if group_list[node] != node:
        group_list[node] = Find(group_list, group_list[node])
    
    return group_list[node]


# Union
def Union(group_list, n1, n2):
    # 그룹
    g1 = Find(group_list, n1)
    g2 = Find(group_list, n2)
    # 같은 그룹이면
    if g1 == g2:
        # 병합하지 않음
        return False, group_list

    # 작은 쪽으로 병합
    if g1 < g2:
        group_list[g2] = g1
    
    else:
        group_list[g1] = g2

    return True, group_list


def solution(N, M, edges):
    # 피로도가 가장 낮은 그룹 / 피로도
    min_group = list(range(N+1))
    min_fat = 0
    # 정렬 / 낮은 그룹
    min_edges = sorted(edges, key=lambda x:-x[-1])
    # 낮은 그룹 병합
    for a, b, c in min_edges:
        # 낮은
        min_state, min_group = Union(min_group, a, b)
        # 병합 되었으면
        if min_state:
            # 피로도 더해주기
            min_fat += 1-c

    # 피로도가 가장 높은 그룹 / 피로도
    max_group = list(range(N+1))
    max_fat = 0
    # 정렬 / 높은 그룹
    max_edges = sorted(edges, key=lambda x:x[-1])
    # 높은 그룹 병합
    for a, b, c in max_edges:
        # 높은
        max_state, max_group = Union(max_group, a, b)
        # 병합 되었으면
        if max_state:
            # 피로도 더해주기
            max_fat += 1-c
    
    # 빼주기
    print(max_fat**2 - min_fat**2)


# 입력
N, M = map(int, input().strip().split())
edges = [list(map(int, input().strip().split())) for _ in range(M+1)]

solution(N, M, edges)