[알고리즘] 탐욕법 - 크루스칼 알고리즘 코드 (Greedy Approach - Kruskal's Algorithm Code)

문제)

Description

교재와 강의자료를 참고하여 Algorithm 4.2 크루스칼 알고리즘의 구현을 완성하시오.

그래프의 입력은 간선의 집합으로 주어지며,

출력값은 크루스칼 알고리즘의 각 단계별로

F 집합에 추가되는 간선을 순서대로 출력한다.

단, 가중치는 unique하며, 동일한 가중치를 가진 간선은 없다고 가정해도 된다.

Input

첫 줄에 정점의 개수 n, 간선의 개수 m이 주어진다.

두 번째 줄부터 다음과 같은 형태로 m 개의 간선이 주어진다.

u v w

u와 v는 정점의 번호이고, 1부터 n까지의 자연수로 표시한다.

w는 간선 <u, v>의 가중치이며, 양의 정수 값으로 주어진다.

Output

크루스칼 알고리즘의 초기 단계부터 종료 단계까지(n-1 times)

F 집합에 포함되는 간선을 순서대로 출력한다.

u v w

Sample Input 1

5 7
1 2 1
1 3 3
2 3 7
2 4 6
3 4 4
3 5 2
4 5 5

Sample Output 1

1 2 1
3 5 2
1 3 3
3 4 4

 

코드)

import sys
from collections import deque
input = sys.stdin.readline


def find(i):
    j = s = i       # 기존의 i를 미리 저장해놓음
    while dset[i] != i:
        i = dset[i]     # 집합의 헤드를 찾으러 감

    while dset[j] != i:     # 이렇게 함으로써 헤드의 부분 집합들이 모두 헤드를 가리키게 만듬
        s = j
        j = dset[j]
        dset[s] = i

    return i


def equal(p, q):
    return True if p == q else False


def union(p, q):
    # p와 q의 집합을 합침. 이때 작은 애가 큰 애를 가리키게 만들어야하고, 작은 애의 집합들이 모두 큰 애의 헤드를 가리키게 만들어야 함
    num_p = num_q = 0
    for i in range(1, vertex_num):      # 어느 집합이 더 작은지 알아냄
        if dset[i] == p:
            num_p += 1
        elif dset[i] == q:
            num_q += 1

    if num_p > num_q:       # 작은 집합을 큰 집합에 흡수 시킴
        dset[q] = p

        for i in range(1, vertex_num):
            if dset[i] == q:
                dset[i] = p

    else:
        dset[p] = q

        for i in range(1, vertex_num):
            if dset[i] == p:
                dset[i] = q
      


def func():
    for i in range(1, vertex_num + 1):
        dset[i] = i     # 각 vertex마다 서로소 집합을 만듬, initial하는 과정

    sorted_edge = deque(sorted(edge, key = lambda x : x[2]))       # weight를 오름차순으로 정렬함

    while len(F) < vertex_num - 1:  # 1부터 vertex_num - 1까지 수행함

        e = sorted_edge.popleft()       # 가중치가 작은 애를 뽑아냄
        i, j = e[0], e[1]       # vertex들을 뽑아냄
        p = find(i)
        q = find(j)
        if not equal(p, q):
            union(p, q)
            F.append(e)
            
            print(*e, sep=' ')    


vertex_num, edge_num = map(int, input().split())        # 정점과 간선의 개수를 입력 받음
edge = deque()
for _ in range(edge_num):
    edge.append(tuple(map(int, input().split())))        # 간선의 리스트를 만듬

F = deque()      # 최소 비용 신장 트리의 간선의 집합
dset = [0] * (vertex_num + 1)       # 서로소 집합을 만듬


func()