Algorithm

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

salmon16 2024. 7. 14. 12:55

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

 

풀이 방법

kruskal 알고리즘을 사용해서 풀이할 수 있었다.

kruskal 알고리즘을 이용하기 위해 union find 알고리즘을 구현해야 한다.

입력을 받을 때 간선(edges)을 기준으로 입력을 받는다.

그 후 edges의 weight가 작은 것으로 정렬을 한다.

weight가 작은 간선 중  간선에 연결된 두 노드의 부모가 같지 않다면 해당 간선을 추가한다.

간선이 V-1개가 될 때까지 반복한다.

 

V , E = map(int, input().split())
parents = [i for i in range(V+1)] ## 부모 배열 생성
edges = []
for i in range(E): ## 입력 받기 
    a, b, w = map(int, input().split())
    edges.append([a, b, w])

def union(a, b): ## 두 개의 부모 중 작은 부모로 합친다.
    parent_a = find(a) ## 부모 구하기
    parent_b = find(b)
    if parent_a < parent_b: ## 작은 부모 구하기
        parents[parent_b] = parent_a ## 부모의 부모를 변경해야 모든 자식 부모가 변경 된다.
    else:
        parents[parent_a] = parent_b

def find(a):
    if parents[a] == a:
        return a
    parents[a] = find(parents[a])
    return parents[a]

edges.sort(key = lambda x : x[2]) ## weight가 낮은 순으로 정렬하기
def mst():
    ret = 0
    cnt = 0
    for a, b, w in edges:
        if find(a) != find(b):
            union(a, b)
            ret += w
            cnt += 1
        if cnt == V-1: ## V-1개 등록되면 종료
            break
    return ret

print(mst())