백준 1647 [도시 분할 계획] 파이썬

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

 

최소 스패닝 트리(MST)를 이용하면 된다는 사실을 알면 쉬운데, 그걸 알아차리기 어려웠다..

 

# 백준 1647 도시 분할 계획
import sys
input = sys.stdin.readline

# N: 집의 개수, M: 길의 개수
N, M = map(int, input().split())
routes = []

# 모든 길 입력받기 (집1, 집2, 비용)
for _ in range(M):
    routes.append(list(map(int, input().split())))

# 비용 기준으로 정렬 (pop()으로 꺼내면서 쓰기 위해 내림차순 정렬)
# 나중에 pop() 하면 비용이 작은 간선부터 나오게 됨
routes.sort(key=lambda x: -x[2])

# 부모 노드 배열 (유니온 파인드용)
parent = [i for i in range(N+1)]

# 선택된 간선들의 비용을 저장
ans = []

# find 함수 (경로 압축)
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

# union 함수 (두 집합을 합침)
def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parent[a] = b
        return True   # 합쳐졌음을 표시
    return False      # 이미 같은 집합이면 합치지 않음

# 크루스칼 알고리즘 실행
while routes:
    A, B, C = map(int, routes.pop())  # 비용이 작은 간선부터 꺼냄
    if union(A, B):                   # 사이클이 생기지 않으면
        ans.append(C)                 # 선택된 간선 비용 저장

# 모든 간선이 연결되면 트리가 만들어지는데,
# 이 문제에서는 도시를 2개로 분할해야 하므로 가장 큰 간선 하나 제거
ans.sort()   # 오름차순 정렬
ans.pop()    # 가장 큰 비용의 간선 제거

# 남은 간선 비용의 합 출력
print(sum(ans))