[프로그래머스] 섬 연결하기

문제 요약

n개의 섬과 섬 사이에 다리를 건설하는 비용이 주어진다.

직접 연결되어 있지 않더라도 다른 섬을 거쳐 이동할 수 있으면 서로 통행 가능한 것으로 본다.

모든 섬이 서로 통행할 수 있도록 다리를 건설할 때 필요한 최소 비용을 구해야 한다.

핵심 아이디어

모든 섬을 연결하면서 전체 비용을 최소로 만들어야 하므로 최소 신장 트리를 구하는 문제다.

최소 신장 트리란 다음 조건을 만족하는 그래프를 의미한다.

모든 정점을 연결한다.
사이클이 없다.
간선 비용의 합이 최소다.

이 문제에서는 크루스칼 알고리즘을 사용할 수 있다.

크루스칼 알고리즘의 진행 과정은 다음과 같다.

모든 다리를 비용이 낮은 순서로 정렬
가장 저렴한 다리부터 확인
두 섬이 아직 연결되어 있지 않다면 다리 선택
이미 연결된 섬이라면 사이클이 생기므로 선택하지 않음
다리를 n - 1개 선택하면 종료

두 섬이 이미 같은 연결 집합에 포함되어 있는지는 Union-Find 자료구조로 확인한다.

최소 신장 트리

섬이 n개일 때 모든 섬을 사이클 없이 연결하기 위해 필요한 다리의 수는 항상 n - 1개다.

예를 들어 섬이 4개라면 최소 신장 트리는 다리 3개로 구성된다.

섬 0 -- 섬 1 -- 섬 2 -- 섬 3

다리를 3개보다 적게 선택하면 모든 섬을 연결할 수 없다.

반대로 이미 모든 섬이 연결된 상태에서 다리를 추가하면 사이클이 만들어진다.

따라서 비용이 낮은 다리부터 사이클이 생기지 않도록 n - 1개를 선택하면 최소 비용을 구할 수 있다.

Union-Find

Union-Find는 여러 원소가 서로 같은 집합에 속하는지 빠르게 확인하고, 서로 다른 두 집합을 합치는 자료구조다.

다음 두 연산을 사용한다.

find: 원소가 속한 집합의 대표 원소를 찾는다.
union: 서로 다른 두 집합을 하나로 합친다.

두 섬의 대표 원소가 같다면 이미 연결된 상태다.

find(island_a) == find(island_b)

이때 다리를 추가하면 사이클이 생기므로 해당 다리를 선택하지 않는다.

대표 원소가 다르면 두 섬을 연결해도 사이클이 생기지 않는다.

find(island_a) != find(island_b)

따라서 다리 비용을 더하고 두 집합을 합친다.

풀이 과정

1. 부모 배열 초기화

처음에는 각 섬이 자기 자신만 포함하는 독립된 집합이다.

parent = list(range(n))

예를 들어 n = 4라면 다음과 같다.

parent = [0, 1, 2, 3]

2. find 함수 구현

섬이 속한 집합의 대표 원소를 찾는다.

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])

    return parent[x]

대표 원소를 찾는 과정에서 부모 값을 바로 대표 원소로 변경한다.

이를 경로 압축이라고 하며, 이후 같은 집합을 다시 탐색할 때 더 빠르게 대표 원소를 찾을 수 있다.

3. union 함수 구현

두 섬의 대표 원소를 찾은 뒤 서로 다른 집합이라면 하나로 합친다.

def union(a, b):
    root_a = find(a)
    root_b = find(b)

    if root_a == root_b:
        return False

    if root_a < root_b:
        parent[root_b] = root_a
    else:
        parent[root_a] = root_b

    return True

이미 같은 집합이면 False를 반환하고 다리를 선택하지 않는다.

서로 다른 집합을 합쳤다면 True를 반환한다.

4. 다리를 비용 순으로 정렬

costs의 각 원소는 다음과 같은 형태다.

[섬 A, 섬 B, 건설 비용]

비용은 세 번째 원소이므로 다음과 같이 정렬한다.

costs.sort(key=lambda bridge: bridge[2])

5. 비용이 낮은 다리부터 선택

정렬된 다리를 순서대로 확인한다.

for island_a, island_b, cost in costs:

두 섬이 서로 다른 집합에 속해 있다면 다리를 선택한다.

if union(island_a, island_b):
    answer += cost
    bridge_count += 1

6. 다리를 n - 1개 선택하면 종료

최소 신장 트리는 항상 n - 1개의 간선을 가진다.

if bridge_count == n - 1:
    break

필요한 다리를 모두 선택했다면 남은 다리는 확인하지 않아도 된다.

Python 코드

def solution(n, costs):
    parent = list(range(n))

    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])

        return parent[x]

    def union(a, b):
        root_a = find(a)
        root_b = find(b)

        # 이미 같은 집합이면 다리를 추가할 수 없다.
        if root_a == root_b:
            return False

        # 서로 다른 두 집합을 하나로 합친다.
        if root_a < root_b:
            parent[root_b] = root_a
        else:
            parent[root_a] = root_b

        return True

    # 건설 비용이 낮은 다리부터 확인한다.
    costs.sort(key=lambda bridge: bridge[2])

    answer = 0
    bridge_count = 0

    for island_a, island_b, cost in costs:
        # 두 섬을 연결해도 사이클이 생기지 않는 경우
        if union(island_a, island_b):
            answer += cost
            bridge_count += 1

            # 최소 신장 트리가 완성되면 종료한다.
            if bridge_count == n - 1:
                break

    return answer

코드 설명

부모 배열

parent = list(range(n))

parent[i]는 i번 섬이 속한 집합의 부모를 나타낸다.

처음에는 모든 섬이 서로 연결되지 않았으므로 각 섬의 부모는 자기 자신이다.

경로 압축

if parent[x] != x:
    parent[x] = find(parent[x])

대표 원소를 찾는 과정에서 거쳐 간 모든 섬의 부모를 대표 원소로 변경한다.

예를 들어 부모 관계가 다음과 같다고 하자.

3 -> 2 -> 1 -> 0

find(3)을 실행한 뒤에는 다음과 같이 변경된다.

3 -> 0
2 -> 0
1 -> 0

이후에는 대표 원소를 훨씬 빠르게 찾을 수 있다.

사이클 판별

if root_a == root_b:
    return False

두 섬의 대표 원소가 같다는 것은 이미 다른 다리들을 통해 연결되어 있다는 의미다.

이 상태에서 두 섬 사이에 다리를 추가하면 사이클이 만들어지므로 선택하지 않는다.

비용 정렬

costs.sort(key=lambda bridge: bridge[2])

크루스칼 알고리즘은 가장 저렴한 간선부터 선택한다.

다만 비용이 낮더라도 사이클을 만드는 다리는 제외한다.

종료 조건

if bridge_count == n - 1:
    break

섬이 n개라면 모든 섬을 연결하기 위해 필요한 다리의 수는 n - 1개다.

해당 개수만큼 다리를 선택한 순간 최소 신장 트리가 완성된다.

예시

다음과 같은 입력을 살펴보자.

n = 4
costs = [
    [0, 1, 1],
    [0, 2, 2],
    [1, 2, 5],
    [1, 3, 1],
    [2, 3, 8]
]

비용을 기준으로 정렬하면 다음과 같다.

[0, 1, 1]
[1, 3, 1]
[0, 2, 2]
[1, 2, 5]
[2, 3, 8]

비용이 낮은 다리부터 선택한다.

0 - 1 연결: 비용 1
1 - 3 연결: 비용 1
0 - 2 연결: 비용 2

3개의 다리로 4개의 섬이 모두 연결된다.

최소 비용 = 1 + 1 + 2 = 4

시간 복잡도

다리의 개수를 E라고 하자.

다리를 비용순으로 정렬하는 데 다음 시간이 필요하다.

O(E log E)

Union-Find 연산은 경로 압축을 사용하므로 매우 빠르게 동작한다.

전체 시간 복잡도는 정렬이 지배하므로 다음과 같다.

O(E log E)

공간 복잡도

각 섬의 부모 정보를 저장하는 parent 배열을 사용한다.

O(n)

단, costs.sort()가 사용하는 정렬 내부 공간은 별도로 볼 수 있다.

정리

이 문제는 모든 섬을 최소 비용으로 연결하는 최소 신장 트리 문제다.

풀이 흐름은 다음과 같다.

모든 다리를 비용순으로 정렬
가장 저렴한 다리부터 확인
Union-Find로 사이클 발생 여부 검사
사이클이 없다면 다리 선택
n - 1개의 다리를 선택하면 종료

최소 신장 트리 문제임을 파악하고 크루스칼 알고리즘과 Union-Find를 함께 사용하는 것이 핵심이다.