[프로그래머스 Lv.3] 섬 연결하기 (Python)

문제 링크: https://school.programmers.co.kr/learn/courses/30/lessons/42861

프로그래머스

풀이

이전에 최소 신장 트리를 그리는 문제에서 사용하였던 Union-find을 통해 구현한 크루스칼 알고리즘을 사용하여 풀이하였다.

 

우선 Union-find를 사용하기 위해 관련 함수들을 구현한다. 이후 크루스칼 알고리즘을 사용하여 최소 신장 트리를 그리는 작업을 수행하면 된다.

먼저 각 섬 간의 건설 비용이 낮은 순서대로 정렬하여, 결과적으로 가장 최저의 비용이 나올 수 있도록 한다.

다음으로 정렬된 건설 비용들을 순회하며, Union-find를 사용해서 최소 신장 트리를 그린다. 이렇게 완성된 최소 신장 트리는 곧 최소 가격으로 모든 섬을 연결하는 케이스와 동일하다. 

def solution(n, costs):
    parents = [i for i in range(n)] # 자기 자신의 루트를 자기 자신으로 초기화
    
    def find(x): # find 메서드 구현
        if parents[x] == x: # 연결되지 않았다면
            return x # 자기 자신을 반환
        parents[x] = find(parents[x]) # 연결된 부모가 있다면, 루트 노드를 찾아감
        return parents[x] # 루트 노드 반환
    
    def union(x, y): # union 메서드 구현
        x = find(x) # x와 연결된 루트를 찾음
        y = find(y) # y와 연결된 루트를 찾음
        
        if x == y: # x와 y의 루트가 같다면, 이미 연결된 것이므로
            return # 즉시 종료
        
        parents[x] = y # 루트가 다르다면 x의 루트를 y로 설정
        
    costs.sort(key=lambda x: x[2]) # cost를 기준으로 오름차순 정렬
    answer = 0
    for cost in costs: # 작은 cost를 가지는 다리부터 순회
        x, y, c = cost
        
        if find(x) == find(y): # 루트가 같다면, 이미 연결된 최소 값에 해당하는 결과가 있으므로
            continue # 건너뜀
        
        union(x, y) # x와 y가 속한 트리를 연결
        answer += c # 연결 비용을 합치기

    return answer