[Programmers Lv.3] 섬 연결하기 - Python

n개의 섬 사이에 다리를 건설하는 비용(costs)이 주어질 때, 최소의 비용으로 모든 섬이 서로 통행 가능하도록 만들 때 필요한 최소 비용을 return 하도록 solution을 완성하세요.
다리를 여러 번 건너더라도, 도달할 수만 있으면 통행 가능하다고 봅니다. 예를 들어 A 섬과 B 섬 사이에 다리가 있고, B 섬과 C 섬 사이에 다리가 있으면 A 섬과 C 섬은 서로 통행 가능합니다.

예시

 

각 섬 사이를 이동할 때 드는 비용이 최소가 될 필요 없이 전체 비용을 최소화하면 되는 문제다. 그렇다면 사이클이 없는 트리 형태로 건설할 때가 최소 비용일 것이다.

MST를 구하는 알고리즘은 여러 가지가 있지만 나는 프림 알고리즘을 사용했다.

 

프림 알고리즘은 그리디 알고리즘의 일종으로 과정은 다음과 같다.

• 임의의 노드를 선택하여 비어있는 트리 T에 포함시킨다.
• T에 있는 노드와 없는 노드 사이의 간선 중, 가중치가 최소인 간선을 찾는다.
• 찾은 간선과 노드를 T에 포함시킨다.
• 모든 노드가 T에 포함될 때까지 반복한다.

 

매번 최소 비용을 찾지 않기 위해 costs 리스트를 미리 정렬하고 시작했다.

트리에 포함시킨 노드를 저장하기 위해 set을 사용했다. 최종 비용만 계산하면 되기 때문에 간선 정보는 저장하지 않았다.

set을 사용한 이유는 코드를 간결하게 만들기 위해서였다. 입력되는 비용 리스트의 형식이 [노드1, 노드2, 비용] 형태인데, 노드1만 트리에 포함된 경우 / 노드2만 트리에 포함된 경우로 나눠서 처리하면 코드가 길어진다. 현재 검토하고 있는 연결 정보에서 트리에 포함된 노드가 하나라면 두 노드 다 트리에 추가하고 최소 비용을 업데이트한다. 중복 노드는 자동으로 처리되어 코드를 추가할 필요가 없다. 

 

def solution(n, costs):
    answer = 0
    costs.sort(key = lambda x:x[2])
    tree = set([0])
    while len(tree) < n: 
        for c in costs:
            if len(set(c[:2]) & tree) == 1:
                tree.update(c[:2])
                answer += c[2]
                break
    return answer

---

얼마 전에 프림 알고리즘과 크루스칼 알고리즘을 공부한 덕분에 문제를 보자마자 MST를 구하는 문제라는걸 캐치할 수 있었다. 전에는 무작정 많이 풀어보면 실력도 는다는 생각이었는데, 이론을 먼저 공부하니까 빨리 풀 수 있는 문제의 범위가 확실히 늘어난다.