개발 · 컴퓨터공학/알고리즘

백준 1197 최소 스패닝 트리 파이썬 문제풀이 (python 최소 스패닝 트리)

2024. 11. 20. 11:43
728x90
반응형

 

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

최소 스패닝 트리(MST)는 그래프가 주어졌을 때, 

높은 가중치보단 가중치가 낮은 edge로 구성되는 최소 비용의 트리를 구하는 문제이다. 

 

최소 스패닝 트리에는 두 가지 방법이 있다.

하나는 다익스트라를 이용하는 프림(prim) 알고리즘.

하나는 크루스칼 알고리즘이다.

 

이를 이용해서 최소 스패닝 트리를 구하는 문제를 풀어본다. 

 

정답 코드 (프림 알고리즘)

import heapq

# prim
N,M = map(int, input().split())

graph = [[] for _ in range(N+1)] # 노드의 개수 + 1
visited = [0 for _ in range(N+1)]

for _ in range(M):
    A,B,C = map(int, input().split())

    graph[A].append([C,B])
    graph[B].append([C,A])

# 다익스트라!

answer = 0
cnt = 0
q = [[0,1]] # 1에서 출발 가중치 없이

while q: # q가 아무것도 없어질 때까지

    if cnt == N:
        break

    weight, node = heapq.heappop(q) # 최소비용을 꺼내준다.

    if visited[node] == 0:
        visited[node] = 1
        answer += weight
        cnt +=1
        for nxt in graph[node]:
            heapq.heappush(q, nxt)

print(answer)

정답 코드 (크루스칼 알고리즘) 

# 크루스칼

# 1. 모든 링크 한 번에 가져온다.
# 2. 링크를 연결하며 같은 집합으로 만든다
# 3. 이미 같은 집합이면 연결 X

def _find(x):
    while par[x] != x: # 루트가 아니면
        x = par[x]
    return x # parent이면 return

def _union(a,b):
    a = _find(a)
    b = _find(b)

    if a == b:
        return

    if rank[a] < rank[b]:
        par[a] = b
    elif rank[b] < rank[a]:
        par[b] = a
    else:
        par[a] = b
        rank[b] += 1

N,M = map(int, input().split())

link = [list(map(int,input().split())) for _ in range(M)]

link.sort(key=lambda x:x[2]) # 가중치 기준 정렬

# [A,B,C]

par = [i for i in range(1_000_001)]
rank = [0 for _ in range(1_000_001)]

ans = 0

for i in range(M):
    A = link[i][0]
    B = link[i][1]
    weight = link[i][2]

    A = _find(A)
    B = _find(B)

    if A == B:
        continue
    else:
        _union(A,B)
        ans += weight
print(ans)

728x90
반응형