[백준] 최소 스패닝 트리(MST)(1197)

union find 함수(합집합 찾기 함수) 를 쓴 Kruskal's 알고리즘으로 풀어보았다.
여기서 핵심은, 같은 집합에 있는 노드끼리는 '같은 부모'를 가지게끔 세팅하는 것이다. 노드값을 숫자로 받을 경우(자연수), 부모 노드를 항상 max값으로 할지, min값으로 할지 'union'함수 내에 녹여내면 된다.

Prim's도 아래에 구현해보겠다.

두 알고리즘 중 어떤 것을 쓸지는 graph의 간선 개수, 밀도에 따라 달려있다. 간선이 적다면, sort가 필요한 Kruskal을 추천하고, 많으면 Prim을 추천한다.

#Pseuso code
parents = [i for i in range(V+1)] #parents초기화;싸이클 테이블이기도 함
edges = []
sumOfEdges = 0

for _ in range(E):
	a,b,cost = input()
	edges.append((cost, a, b))

edges.sort() # cost 기준으로 오름차순 정렬

#union-find 함수 구현
def find(parents, x):
	if parents[x]!=x:
    	return find(parents, parents[x])
    else:
    	return parents[x]

def union(parents, a, b):
	parents_a = find(parents, a) # a의 부모노드
    parents_b = find(parents, b) # b의 부모노드
    if parents_a < parents_b: #부모 노드를 min으로 두게 해보자
    	parents[parents_b] = parents_a
    else:
    	parents[parents_a] = parents_b

#실행
if find(parents, a) != find(parents, b):
	union(parents, a, b)
    sumOfEdges += cost

find 함수는 '최적화'된 상태이다. else: return parents[x] 가 핵심 구문인데, return x만 할 때보다 메모리적으로 훨씬 우수하다고 볼 수 있다. 자신의 부모노드를 parents에 저장할 것인지, 아니면 부모of부모(of부모,,,)노드를 parents에 저장할 것인지의 차이이다.

아래는 prim's 알고리즘이다.

임의의 노드를 선택 후,
인접 노드들을 minheap에 넣는다.
최소 cost의 노드를 뽑아, 방문하지 않았다면 선택한다.

import heapq
import sys

sys.setrecursionlimit(10**6)
input = sys.stdin.readline

v,e = map(int, sys.stdin.readline().split())
visited = [0 for _ in range(v+1)]
edges = [[] for _ in range(v+1)]

for i in range(e):
    x, y, weight = map(int, sys.stdin.readline().split())
    edges[x].append([weight, y])
    edges[y].append([weight, x])

heap=[]
answer = 0
cnt = 0

heap.append([0,1])

while cnt < v:
    weight, node = heapq.heappop(heap)
    if visited[node] == 0:
        visited[node] = True
        answer += weight
        cnt+=1
        for i in edges[node]:
            heapq.heappush(heap, i)

print(answer)