[프로그래머스] 섬 연결하기 (Python) - 유니온 파인드 알고리즘

문제 설명

https://school.programmers.co.kr/learn/courses/30/lessons/42861

프로그래머스

 

 

 

코드 구현

# 실패 - 사이클을 감지하지 못함

대표 예시는 통과했지만, 막상 제출하고 보니 대부분의 테스트 케이스를 통과하지 못한 저의 첫 번째 풀이입니다. 어디서부터 잘못됐는지 확인하기 위해 문제를 다시 읽어보니, 이 코드로는 사이클을 감지하지 못한다는 점을 깨달았습니다. 처음에는 최소 비용만을 고려한 단순한 트리 문제라고 생각했지만, 본질은 최소 신장 트리(MST)였던 것이죠.

 

  

    
    
    
    
  
def solution(n, costs):
    answer = 0  # 총 건설 비용
    connected = set()  # 연결된 섬
    costs.sort(key=lambda x: x[2])  # costs 건설 비용 측면에서 오름차순 정렬
    
    for n1, n2, cost in costs:
        if len(connected) == n:
            return answer
        if not n1 in connected or not n2 in connected:
            connected.add(n1)
            connected.add(n2)
            answer += cost

 

 

# 정답 풀이

유니온-파인드 알고리즘과 최소 신장 트리 속성을 이용해, 아래처럼 문제를 해결할 수 있습니다. 유니온-파인드 알고리즘에 대한 설명은 아래 게시글을 참고해주세요 :)

 

  

    
    
    
    
  
def find(x, parents):
    if parents[x] != x:
        parents[x] = find(parents[x], parents)
    return parents[x]


def union(r1, r2, parents, rank):
    root1 = find(r1, parents)
    root2 = find(r2, parents)
    
    if rank[root1] < rank[root2]:
        parents[root1] = root2
    elif rank[root2] < rank[root1]:
        parents[root2] = root1
    else:
        parents[root2] = root1
        rank[root1] += 1
    

def solution(n, costs):
    costs.sort(key=lambda x: x[2])  # costs 건설 비용 측면에서 오름차순 정렬
    
    min_cost = 0
    rank = [0]*n
    parents = [i for i in range(n)]
    num_edge = 0
    
    for n1, n2, cost in costs:
        # 최소 신장 트리 속성 이용
        if num_edge == n-1:
            break
        
        # 현재 간선의 두 노드의 루트 노드 찾기
        r1 = find(n1, parents)
        r2 = find(n2, parents)
        
        # 루트 노드가 다르다면 두 집합을 합침
        if r1 != r2:
            union(r1, r2, parents, rank)
            min_cost += cost
            num_edge += 1
            
    return min_cost

 

 

[집합] 개념 정리 (Feat. 유니온-파인드 알고리즘)