[Level 3] 섬 연결하기

programmers.co.kr/learn/courses/30/lessons/42861

코딩테스트 연습 - 섬 연결하기

---

접근

문제를 읽었을 때 n개의 섬을 연결하는 최소신장트리를 구성하는 것을 알 수 있었고 크루스칼 알고리즘을 적용하였다.

 

구현

- 입력받은 costs를 비용을 기준으로 오름차순 정렬한다.

- 이후 자신의 최상위 부모노드를 입력하기 위해 리스트 p를 선언하고 초기에는 자기자신을 가리키도록 한다.

    costs = sorted(costs, key=lambda x: x[2])
    p = [i for i in range(n)]

- 자신의 최상위 부모노드를 찾는 find_set함수를 구현하였다. 만약 자기자신이 아닐 때는 현재 자신의 부모에 대한 최상위 부모 노드를 찾아서 반환한다.

- 이후 두 개의 노드를 연결할 수 있도록 union 함수를 구현하였다.

    def find_set(a):
        if a == p[a]:
            return a
        else:
            b = find_set(p[a])
            return b
    def union(a, b):
        a = find_set(a)
        b = find_set(b)
        if a != b:
            p[b] = a

- 이제 n-1개의 선을 구한다.

- costs는 비용이 적은 것부터 탐색되도록 정렬되어있으며 

- 두 개의 섬인 x와 y가 같은 부모가 아닐 때 두 개의 섬을 연결한다.

- 그리고 해당 비용을 answer에 더하며 선택한 선의 수인 picked를 1증가 시킨다.

- 이후 picekd가 n-1이 되면 위의 과정을 종료한다.

    picked = 0
    for x, y, cost in costs:
        if find_set(x) != find_set(y):
            union(x, y)
            answer += cost
            picked += 1
        if picked == n-1:
            break

---

전체 코드

def solution(n, costs):
    answer = 0
    costs = sorted(costs, key=lambda x: x[2])
    p = [i for i in range(n)]
    def find_set(a):
        if a == p[a]:
            return a
        else:
            b = find_set(p[a])
            return b
    def union(a, b):
        a = find_set(a)
        b = find_set(b)
        if a != b:
            p[b] = a
    
    picked = 0
    for x, y, cost in costs:
        if find_set(x) != find_set(y):
            union(x, y)
            answer += cost
            picked += 1
        if picked == n-1:
            break
    return answer