[Programmers/Python] Greedy - 섬 연결하기

Programmers - 섬 연결하기

문제

n개의 섬 사이에 다리를 건설하는 비용(costs)이 주어질 때, 최소의 비용으로 모든 섬이 서로 통행 가능하도록 만들 때 필요한 최소 비용을 return 하도록 solution을 완성하세요.

다리를 여러 번 건너더라도, 도달할 수만 있으면 통행 가능하다고 봅니다. 예를 들어 A 섬과 B 섬 사이에 다리가 있고, B 섬과 C 섬 사이에 다리가 있으면 A 섬과 C 섬은 서로 통행 가능합니다.

제한사항

• 섬의 개수 n은 1 이상 100 이하입니다.
• costs의 길이는 ((n-1) * n) / 2이하입니다.
• 임의의 i에 대해, costs[i][0] 와 costs[i][1]에는 다리가 연결되는 두 섬의 번호가 들어있고, costs[i][2]에는 이 두 섬을 연결하는 다리를 건설할 때 드는 비용입니다.
• 같은 연결은 두 번 주어지지 않습니다. 또한 순서가 바뀌더라도 같은 연결로 봅니다. 즉 0과 1 사이를 연결하는 비용이 주어졌을 때, 1과 0의 비용이 주어지지 않습니다.
• 모든 섬 사이의 다리 건설 비용이 주어지지 않습니다. 이 경우, 두 섬 사이의 건설이 불가능한 것으로 봅니다.
• 연결할 수 없는 섬은 주어지지 않습니다.

입출력 예

입출력 예 설명

costs를 그림으로 표현하면 다음과 같으며, 이때 초록색 경로로 연결하는 것이 가장 적은 비용으로 모두를 통행할 수 있도록 만드는 방법입니다.

풀이

설계

'최소 신장 트리'를 찾아내는 문제로서, 그리디의 일종인 '크루스칼 알고리즘(Kruskal algorithm)'을 사용해 해결하는 문제이다. 크루스칼 알고리즘 포스트에 전반적인 내용이 있으니 여기서는 간단하게만 정리한다.

1. costs의 리스트를 비용 기준 오름차순으로 정렬한다.
2. 각 노드들의 루트 노드를 기록하기 위한 parent 리스트를 생성하고, 초깃값으로 각 노드들 자신의 값을 넣어 준다.
3. costs의 요소를 하나씩 꺼내 Union-Fide 연산을 활용, 각 노드의 루트 노드를 찾아 parent의 값을 갱신한다.
4. 3과 동시에 해당 노드 간의 사이클 발생 여부를 확인, 사이클이 발생하지 않는 간선의 cost를 answer에 더해 준다.

코드

def solution(n, costs):
    # 간선 정보를 비용을 기준으로 정렬
    costs.sort(key=lambda x: x[2])

    # 부모 노드 테이블 생성 후 초기화
    parent = [0] * (n)
    for i in range(n):
        parent[i] = i

    # find 함수와 union 함수 선언
    def find(parent, x):
        # 자기 자신이 루트 노드가 아닐 경우 루트 노드를 찾을 때까지 재귀적으로 호출
        if parent[x] != x:
            # return find(parent, parent[x])
            parent[x] = find(parent, parent[x])  # 불필요한 재귀를 줄이기 위해 테이블 즉시 갱신
        return parent[x]

    def union(parent, a, b):
        a = find(parent, a)
        b = find(parent, b)
        # 각각의 부모 노드 값을 비교하여 작은 값으로 갱신
        if a < b:
            parent[b] = a
        else:
            parent[a] = b

    # 간선 정보를 순회하면서 크루스칼 알고리즘 수행
    def kruskal(costs, parent):
        tree_cost = 0
        for a, b, cost in costs:
            # 두 노드 간의 사이클 여부 확인
            if find(parent, a) != find(parent, b):
                union(parent, a, b)
                tree_cost += cost
        return tree_cost
    
    answer = kruskal(costs, parent)
    return answer