크루스칼, 프림스, BOJ [최소 스패닝 트리]

문제

문제 보기

다음과 같이 그래프의 간선들이 주어질 때 최소 스패닝 트리의 weight를 구하시오. a,b,c 로 주어질 때 a,b를 잇는 간선의 weight는 c라는 뜻이다.

얻은 개념

문제는 어렵지 않게 풀었는데 Kruskal과 Prim's 알고리즘을 잘 사용해보지 않았어서 이번기회에 정리하고자 한다. 그전에 for else 문을 이중 for 문에서 사용하는 방법이다.

for arr1_idx, arr1_val in enumerate(arr1):
    for arr2_idx, arr2_val in enumerate(arr2):
        if arr1_val == arr2_val:
            print("탈출")
            break
    else:
        continue
    break

알고리즘 공부 링크

각 알고리즘에 대한 설명은 링크로 대신하고 여기서는 링크에 게시된 코드를 설명하겠다. 코드 복붙이 안돼서 걍 줄글로 설명한다.

Kruskal 알고리즘의 핵심은 find-union이다. edge들을 가중치 순서로 정렬한 다음 하나씩 뽑으면서 두 노드의 root가 다르면 합치는 것이다. 모든 edge에 대해서 수행하면 알아서 clique가 생기는 애들은 걸러진다.

Prim's 알고리즘은 bfs + heapq를 사용해서 현재 연결된 노드에서 갈 수 있는 가장 weight가 작은 간선을 계속 선택해서 확장한다. visited 처리를 해서 한번 방문한 node는 다시 방문하지 않게하여 clique르 없앤다. 처음에는 아무 node나 random하게 선택해서 시작한다. 어차피 모두 이어질 것이기 때문이다.

코드

import heapq

v,e = map(int,input().split())
graph = [[] for _ in range(v+1)]
edges = []
visited = [False] * (v+1)

for _ in range(e):
	edges.append(list(map(int,input().split())))
	a,b,c = edges[-1]
	graph[a].append((c,b))
	graph[b].append((c,a))

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

visited_num = 0
weight = 0

q = []
edge = edges[0]

visited[edge[0]] = True
visited[edge[1]] = True
visited_num += 2
weight += edge[2]

for i in graph[edge[0]]:
	heapq.heappush(q,i)
for i in graph[edge[1]]:
	heapq.heappush(q,i)

while visited_num<v:
	cost, node = heapq.heappop(q)
	
	if visited[node]:
		continue
	
	weight += cost
	visited[node] = True
	visited_num += 1
	for i in graph[node]:
		heapq.heappush(q,i)

print(weight)