接続プログラマー-島

問題を解く

Union Findアルゴリズム

プールコード

def solution(n, costs):
    answer = 0
    
    # 특정 원소가 속한 집합을 찾기
    def find_parent(parent, x):
        # 루트가 아니라면 찾을 때까지 재귀 호출
        if parent[x] != x:
            parent[x] = find_parent(parent, parent[x])
        return parent[x]

    # 두 원소가 속한 집합을 찾기
    def union_parent(parent, a, b):
        a = find_parent(parent, a)
        b = find_parent(parent, b)
        if a < b:
            parent[b] = a
        else:
            parent[a] = b
    
    # 부모 테이블
    parent = [0] * n
    # 부모를 자기 자신으로 초기화
    for i in range(n):
        parent[i] = i
    
    # 크루스칼 알고리즘을 위해 비용순으로 정렬
    costs = sorted(costs, key=lambda x:x[2])
    
    # 크루스칼 알고리즘 수행
    # 간선을 하나씩 확인하며
    for edge in costs:
        a, b, cost = edge
        # 사이클이 발생하지 않는 경우에만 집합에 포함
        if find_parent(parent, a) != find_parent(parent, b):
            union_parent(parent, a, b)
            answer += cost
    
    return answer