[백준 1197번] 최소 스패닝 트리 : MST(Minimum Spanning Tree) - 파이썬(Python)

https://www.acmicpc.net/problem/1197

1197번: 최소 스패닝 트리

parent = {}
rank = {}

def make_set(V):
    for v in range(1, V+1):
        parent[v] = v # 부모 정점을 자신으로 초기화한다.
        rank[v] = 0 # rank가 높은 정점이 부모가 된다.
        
def find(v): # root 정점을 반환한다.
    if parent[v] != v:
        parent[v] = find(parent[v])
    return parent[v]

'''
<Union-by-Rank>
높이를 rank로 기억한다.
1. 두 트리의 높이가 다를 때 => 높이가 작은 트리를 높이가 큰 트리에 붙인다.
2. 두 트리의 높이가 같을 때 => 한 쪽 트리의 높이를 1 증가시켜 다른 트리의 root 정점으로 한다.
'''
def union(u, v):
    root1 = find(u)
    root2 = find(v)
    
    if root1 != root2: # root가 다르면 rank가 높은 정점이 부모가 된다.
        if rank[root1] > rank[root2]:
            parent[root2] = root1
        else:
            parent[root1] = root2
            if rank[root1] == rank[root2]:
                rank[root2] += 1 # rank가 같다면 root2를 root1의 부모로 만들었으므로 root2의 rank를 1 올린다.

'''
<Kruskal 알고리즘>
탐욕적인 방법(greedy method)을 이용하여 그래프 내의 모든 정점들을 최소 비용으로 연결
'''
def kruskal(edge_list, V):
    total_cost = 0
    edge_list.sort() # 정점의 리스트를 cost 기준으로 오름차순으로 정렬한다. (탐욕적인 방법을 사용하기 때문.)
    
    make_set(V)
    
    for e in edge_list: # 정렬된 edge_list의 원소를 차례로 가져온다.
        cost, a, b = e
        if find(a) != find(b): # root 정점이 다르면
            union(a, b) # root 정점을 같게 만들어준다.
            total_cost += cost
            
    return total_cost

def main():
    V, E = map(int, input().strip().split(' '))
    edge_list = []
    for _ in range(E):
        a, b, c = map(int, input().strip().split(' '))
        edge_list.append((c, a, b))
        
    print(kruskal(edge_list, V))
    
if __name__ == '__main__':
    main()