[Python] 1197 - 최소 스패닝 트리 (골드4)

1. 문제 설명

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

1197번: 최소 스패닝 트리

설명이 좀 불친절한 감이 있다.

문제를 못 풀 정도로 지장이 생기진 않지만, 예외처리 범위를 생각하느라 괜히 생각이 많아진다.

 

---

2. 아이디어

오랜만에 문제를 봤더니 아이디어가 바로바로 안 떠올랐다.

어차피 골드니까 슬슬 풀고 다음 포스팅으로 넘어가려다가 나중의 나를 위해 정리하고 넘어가기로 했다.

참고로 나는 union-find 알고리즘을 풀었었다.

 

만약 이런 정점과 간선 정보가 주어졌다면 어느 경로로 이동해야 모든 노드를 움질일 때

최소 비용으로 이동할 수 있겠는가? 라는 질문이다.

input으로 (1 2 1), (1 3 3) 이런 식으로 주기 때문에 화살표를 단일로 그렸을 뿐, 쌍방으로 이동할 수 있다고 생각하자.

 

모든 경우를 고려해보았을 때, 1 -> 2 -> 3번 정점 순서로 이동하는 것이 가장 가중치가 적다.

그렇다면 여기서 생각해볼 수 있는 방법은 정점을 이어주기 전에 비용이 적은 순서로 정렬을 하고

그 다음에 union-find로 이어주면 되지 않을까? 라는 생각이다.

 

위의 그림과 같이 정점이 5개와 간선 정보가 5개 주어졌다고 하자.

정점(u) 이동할 정점(v) 비용(w)라고 했을 때, 붉은색 글씨로 써놓은 순서대로 흘러갈 것이다.

 

1번과 4번 정점을 보니 서로 이어져 있지 않으므로 연결해준다.

쭉쭉 이어준다.

 

그런데 정점 2번과 정점 1번을 이으려고 보니 1번 정점은 (1 -> 4 -> 3)이므로 3번 정점으로 이어져있고,

2번 정점은 (2 -> 3)이므로 3번 정점과 이어져 있으므로 이 둘은 이미 연결되어 있다.

그렇다면 이 케이스는 더이상 필요없는 간선 정보가 되므로 넘어간다.

 

마지막 남은 정보를 마저 이어주면 통과할 수 있는 문제였다.

 

다만, 내가 이 당시에 썼던 Union-find 알고리즘이 정점 정보를 순서대로 읽는 문제를 풀 때 쓰던 걸

똑같이 구현해버려서 중간에 불필요한 로직이 하나 들어간다.

0, 4, 3, 5, 3, 5가 아니다.

결과에 영향을 주거나 시간이 많이 들어가는 것도 아니라 단순히 union 함수에 조건이 몇 개 들어간 거라

그냥 알고만 있어도 무방하다.

 

---

3. 코드

import sys

input = sys.stdin.readline

def find(a, node):
    if a == node[a]:
        return a
    return find(node[a], node)

def union(a, b, node):
    a = find(a, node)
    b = find(b, node)

    if a > b:
        node[b] = a
    elif a <= b:
        node[a] = b


def solution():
    v, e = map(int, input().split())
    tree = [i for i in range(v+1)]
    cost = []
    result = 0

    for _ in range(e):
        a, b, c = map(int, input().split())
        cost.append((c, a, b))
    cost.sort()

    for value in cost:
        c, a, b = value
        if find(a, tree) != find(b, tree):
            union(a, b, tree)
            result += c
    print(result)

solution()