[백준/BOJ] gold4 - 1197번 최소 스패닝 트리 (Python)

▶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보다 작거나 같은 데이터만 입력으로 주어진다.

 

▶출력

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

728x90

▶풀이

크루스칼이랑 프림 두 가지 방식으로 풀었다.

두 방식 모두 python3, pypy3 통과하는데 문제가 없었다.

 

Kruskal 알고리즘

def find(x):
    if x != root[x]:
        root[x] = find(root[x])
    return root[x]


v, e = map(int, input().split())
root = [i for i in range(v + 1)]
arr = []
answer = 0

for _ in range(e):
    arr.append(list(map(int, input().split())))

arr.sort(key=lambda x: x[2])

for a, b, c in arr:
    ar = find(a)
    br = find(b)

    if ar != br:
        if ar > br:
            root[ar] = br
        else:
            root[br] = ar
        
        answer += c

print(answer)

Prim 알고리즘

import heapq


v, e = map(int, input().split())
arr = [[] for _ in range(v + 1)]

visit = [0 for _ in range(v + 1)]
heap = [[0, 1]]
answer = 0
cnt = 0

for _ in range(e):
    a, b, c = map(int, input().split())
    arr[a].append([c, b])
    arr[b].append([c, a])

while heap:
    if cnt == v:
        break

    c, b = heapq.heappop(heap)
    if visit[b] == 0:
        visit[b] = 1
        answer += c
        cnt += 1
        for tmp in arr[b]:
            heapq.heappush(heap, tmp)

print(answer)

 

두 가지 방식 모두 상관은 없지만, 개인적으로 union, find함수를 사용하는 크루스칼보다는

함수 없이 바로 코드를 짤 수 있는 프림을 선호한다.