[BOJ/백준] 1197 : 최소 스패닝 트리

Webb ㅣ 2023. 7. 1. 12:32

문제 링크
백준 1197 : 최소 스패닝 트리

문제

그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

아이디어

최소 스패닝 트리 (Minimum Spanning Tree) 를 찾는 알고리즘을 구현하기만 하면 되는 간단한 문제이다.
MST를 찾기 위해 사용되는 알고리즘은 크게 2가지가 있다.

크루스칼(Kruskal) 알고리즘
프림(Prim) 알고리즘

크루스칼 알고리즘은 모든 엣지를 가중치가 작은 순서대로 정렬한 다음에, 추가한 엣지가 사이클을 만들지 않는 경우(=같은 노드를 중복해서 선택하지 않는 경우)에만 추가한다. 시간복잡도는 다음과 같다.

$$
O(E\log V)
$$

반면 프림 알고리즘은 임의의 노드를 하나씩 선택해가면서, 아직 선택되지 않은 노드와 연결되는 엣지중에서 가중치가 가장 작은 엣지를 선택해 나간다. 시간복잡도는 다음과 같다.

$$
O(V^2),\ \text{with priority Q : } O((V+E)\log V)
$$

이 문제를 풀기시작한 이유가 구현 연습이기 때문에 2가지 알고리즘 모두 구현해보고 성능까지 비교해보려고 한다.

알고리즘 : 코드

크루스칼 알고리즘
(*사이클 유무를 확인할 때는 UNION-FIND 사용)

import sys
input = sys.stdin.readline

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]  # 경로압축 (direct-root)

def union(a, b):
    pA = find(a)
    pB = find(b)
    if pA == pB:
        return
    if pA < pB:
        parent[pB] = pA
    else:
        parent[pA] = pB

# 크루스칼
def kruskal():
    weight = 0
    edges.sort()
    for w, u, v in edges:
        if find(u) == find(v):
            continue
        weight += w
        union(u, v)
    return weight

V, E = map(int, input().split())
parent = list(range(V + 1))
edges = []
for _ in range(E):
    u, v, w = map(int, input().split())
    edges.append([w, u, v])

result = kruskal()
print(result)

프림 알고리즘

import sys
import heapq

input = sys.stdin.readline

# 프림
def prim(n, w):
    totalWeight = 0
    queue = [[w, n]]  # heap 특징을 쓰기 위해 weight를 앞으로 바꿔줘야함s
    while queue:
        weight, node = heapq.heappop(queue)
        if not visited[node]:
            visited[node] = True
            totalWeight += weight
            for nNode, nWeight in graph[node]:
                heapq.heappush(queue, [nWeight, nNode])
    return totalWeight

V, E = map(int, input().split())
graph = [[] for _ in range(V + 1)]
visited = [False for _ in range(V + 1)]
for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append([v, w])
    graph[v].append([u, w])

result = prim(1, 0)
print(result)

두 알고리즘을 사용했을 때 메모리와 수행시간을 다음과 같이 차이난다. (위 : 크루스칼, 아래 : 프림)

물론 완벽히 최적화 된 코드가 아니긴 하지만 크루스칼 알고리즘이 프림에 비해 훨씬 더 효율적으로 동작한다는 것을 알 수 있었다.

회고

이번에 썼던 short paper에서 네트워크 라우팅을 설명할 때 MST를 도입했었는데 개념만 알고있고 실제로 구현해본적이 없었던 것 같아서 해당 문제를 풀어보았다.
역시 개념만 아는 것과 실제 코드로 구현하는 것 사이에는 많은 간극이 존재하는 듯 하다.

크루스칼 짤 때 FIND 최적화 구현하는게 x를 parent[x]로 수정하는 것 만으로도 할 수 있다는게 신기했고 heap 자료구조를 사용할 때 아무생각없이 그냥 node, weight순으로 넣다가 뒤늦게 문제를 알았었다. 자료구조에 조금 더 친숙해져야 할 것 같다.

내 꿈은 척척학사