백준 1197 최소 스패닝 트리 문제 풀이

문제 알아보기

가장 기본적인 최소 비용 스패닝 트리 문제입니다. 제시된 그래프에서 최소 스패닝 트리를 구성할 수 있는 간선만 선택하여 간선 가중치의 합을 구합니다.

풀어보기

프림 알고리즘

대표적인 최소 비용 스패닝 트리 알고리즘인 프림 알고리즘을 사용합니다.

1. 간선의 가중치를 기준으로 모든 간선 정보를 정렬.
2. 가중치가 가장 낮은 간선을 선택.
3. 남아있는 간선들 중 사이클이 형성되지 않으면서 가중치가 가장 낮은 간선을 선택.
4. 3 반복

프림 알고리즘은 다익스트라 알고리즘과 매우 유사하고, 알고리즘 자체가 간단하여 쉽게 구현할 수 있습니다.

위와 같이 복합적인 값을 일정한 기준으로 정렬할 때, 개인적으로 가장 용이하게 사용하는 방법은 클래스를 선언하여 비교자를 재정의하고 정렬하는 것입니다.

이 문제의 경우 간선에 대한 클래스를 선언하고 비교자가 가중치를 사용하여 두 객체를 사용하도록 재정의하여 가중치를 기준으로 정렬하도록 할 수 있습니다.

class Conn:
    def __init__(self, a: int, b: int, cost: int):
        self.a = a
        self.b = b
        self.cost = cost

    def __eq__(self, other) -> bool:
        return self.cost == other.cost

    def __lt__(self, other) -> bool:
        return self.cost < other.cost

    def __le__(self, other) -> bool:
        return self < other or self == other

위와 같이 정의하면 Conn 객체는 가중치(cost)를 기준으로 정렬될 것입니다.

conns: list[Conn] = []
...
conns.sort()

사이클 형성 여부는 유니온-파인드를 사용하여 판단합니다. 두 정점의 최상위 부모가 같다면 이미 어떠한 형태로든 두 정점은 연결되어있다는 것이므로, 두 정점을 더 연결하면 사이클이 발생할 것입니다.

if find(conn.a) == find(conn.b):
    continue
else:
    merge(conn.a, conn.b)

답안 구현

1197 Python 답안