[백준] 학교 탐방하기 (python)

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

 

풀이 방법

union-find알고리즘을 통해 해결할 수 있다. 

기존의 알고리즘과 다른 점은 0에서 1로 가는 간선은 무조건 포함되어야 한다는 점이다. 

또 가장 힘든 경우와 가장 덜 힘든 경우를 계산하기 위해 mst알고리즘을 2번 사용했다.

주의: mst알고리즘을 사용하기 전 parents를 초기화해주어야 한다.

 

 

N, M = map(int, input().split())
edges = []
parents = [ i for i in range(N+1)]
a, b, c = map(int, input().split()) ## 0~1 엣지 무조건 추가되어야 할 edge
first_weight = 1 - c ## 첫 번째 edge의 weight 오르막인지 내리막인지 

for i in range(M): ## 엣지 추가
    a, b, c = map(int, input().split())
    edges.append([a, b, c])

def union(a, b): ## union 알고리즘
    a = find(a)
    b = find(b)
    if a < b:
        parents[b] = a
    else: 
        parents[a] = b

def find(a): ## find알고리즘
    if parents[a] == a:
        return a
    
    parents[a] = find(parents[a])
    return parents[a]

def mst(): 
    cnt = 0 ## 0~1엣지 포함
    weight = first_weight ## 0~1 edge는 무조건 포함
    for i in range(M):
        if cnt == N-1: ## 모든 엣지를 포함하면 break
            break
        a, b, c = edges[i]
        if find(a) != find(b):
            union(a, b) 
            cnt += 1 ## 카운드 + 1
            if c == 0: ## 오르막인 경우
                weight += 1 ## weight 1추가
    return weight ** 2 ## weight 제곱 
## 내리 막 위주로 되게 정렬 0이 오르막 1이 내리막
## 가장 비용이 작은 경우
edges.sort(key = lambda x : -x[2])
first = mst()
## 가장 비용이 큰 경우
## 부모 초기화
parents = [ i for i in range(N+1)]
edges.sort(key = lambda x : x[2])
second = mst()
print(second - first)