[백준] 학교 탐방하기

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

13418번: 학교 탐방하기

MST 문제입니다. 최악, 최선의 피로도를 구해야하는데 최악의 경우에는 내리막길(0), 최선의 경우에는 오르막길(1)이 먼저나오도록 간선들을 정렬해서 크루스칼을 돌려주면 됩니다.

 

import sys

sys.stdin = open("input.txt", "r")
input = sys.stdin.readline

n, m = tuple(map(int, input().split()))

edges = [
    tuple(map(int, input().split()))
    for _ in range(m + 1)
]

edges.sort(key=lambda x: x[2])


parent = []

def find(p):
    if parent[p] == p:
        return parent[p]

    parent[p] = find(parent[p])
    return parent[p]

def kruskal():
    global parent
    parent = [
        x for x in range(n + 1)
    ]

    ret = 0

    for i in range(m):
        a, b, w = edges[i]
        pa = find(a)
        pb = find(b)

        if pa == pb:
            continue

        ret += (1 if w == 0 else 0)

        if pa < pb:
            parent[pb] = pa
        else:
            parent[pa] = pb

    return ret

max_dist = kruskal()

edges.sort(key=lambda x: -x[2])

min_dist = kruskal()

print(max_dist ** 2 - min_dist ** 2)