백준(BOJ) 13418 학교 탐방하기(Python)

문제 내용.

 

MST 문제.

최대 스패닝 트리와 최소 스패닝 트리를 각각 구한 다음,

각각의 코스트의 제곱의 차를 출력하면 된다.

 

코드는 다음과 같다.

 

# -*- coding: utf-8 -*-
import sys
from collections import deque
import heapq
import bisect
import math
from itertools import product
from itertools import combinations
"""
from itertools import combinations
from itertools import combinations_with_replacement
from itertools import permutations
import copy
"""
#input=sys.stdin.readline
#print=sys.stdout.write
#sys.setrecursionlimit(100000000)


#Union-Find
def find(x):
    if x!=parent[x]:
        parent[x]=find(parent[x])
    return parent[x]

def union(x,y):
    x=find(x)
    y=find(y)
    x,y=min(x,y),max(x,y)
    if x!=y:
        parent[y]=x


n,m=map(int,input().split())
L1=[]
L2=[]

for i in range(m+1):
    a,b,c=map(int,input().split())
    heapq.heappush(L1,[c^1,a,b])
    heapq.heappush(L2,[-(c^1),a,b])

parent=[i for i in range(n+1)]

count=0
ans1=0
while count<n:
    c,a,b=heapq.heappop(L1)
    if find(a)!=find(b):
        count+=1
        union(a,b)
        ans1+=c
parent=[i for i in range(n+1)]

count=0
ans2=0
while count<n:
    c,a,b=heapq.heappop(L2)
    if find(a)!=find(b):
        count+=1
        union(a,b)
        ans2+=c
print(ans2*ans2-ans1*ans1)

의욕이 없다 요즘.