파이썬으로 최소 비용으로 모든 도시 연결하기: 크루스칼 알고리즘 풀이

문제 설명

1부터 N까지 번호가 매겨진 N개의 도시가 있다고 가정해 봅시다. connections 배열에는 여러 개의 연결 정보가 담겨 있으며, 각 연결은 [city1, city2, cost] 형태로 city1과 city2를 직접 연결하는 데 드는 비용을 나타냅니다. 목표는 모든 도시 쌍 사이에 경로(길이 1인 직접 연결도 포함)가 존재하도록 도시들을 잇는 것이며, 이때 총 비용은 선택한 연결들의 비용 합계가 됩니다. 모든 도시를 연결하는 것이 불가능하다면 -1을 반환해야 합니다.

예를 들어 그래프가 다음과 같다고 해보겠습니다.

이 경우 출력은 6입니다. 두 개의 연결만 선택해도 모든 도시가 하나로 이어지므로, 비용이 가장 낮은 간선들(2와 [1, 5])을 선택하는 것이 최적입니다.

접근 방법: 크루스칼 알고리즘 + 유니온-파인드

이 문제는 대표적인 최소 신장 트리(MST, Minimum Spanning Tree) 문제입니다. 간선을 비용 순으로 정렬한 뒤 사이클을 만들지 않는 간선부터 차례로 선택하는 크루스칼(Kruskal) 알고리즘과, 두 노드가 같은 집합에 속했는지 빠르게 판별하는 유니온-파인드(Union-Find, 서로소 집합) 자료구조를 조합하면 효율적으로 해결할 수 있습니다.

풀이 절차는 다음과 같습니다.

find() 메서드를 정의합니다. 인자로 x를 받습니다.
parent[x]가 -1이면 x를 그대로 반환합니다.
그렇지 않으면 parent[x] := find(parent[x])로 재귀 호출해 루트를 찾고(경로 압축), 그 결과를 반환합니다.
union() 메서드를 정의합니다. 인자로 x와 y를 받습니다.
parent_x := find(x), parent_y := find(y)를 구합니다.
두 값이 같으면 이미 같은 집합이므로 그대로 반환하고, 그렇지 않으면 parent[parent_y] := parent_x로 두 집합을 하나로 합칩니다.
메인 메서드에서는 n과 connections를 입력받습니다.
크기가 n+1인 parent 배열을 -1로 초기화하고, disjoint := n - 1(아직 연결해야 할 남은 도시 수), cost := 0으로 설정합니다.
c := connections를 세 번째 요소(비용)를 기준으로 오름차순 정렬한 리스트로 만듭니다.
i := 0으로 초기화합니다.
i가 c의 길이보다 작고 disjoint가 0이 아닌 동안 반복합니다.
- x := c[i][0], y := c[i][1]
- find(x)와 find(y)가 다르면 disjoint를 1 감소시키고, cost에 c[i][2]를 더한 뒤 union(x, y)를 수행합니다.
- i를 1 증가시킵니다.
반복이 끝난 후 disjoint가 0이면 cost를, 그렇지 않으면 -1을 반환합니다.

예제 코드 (파이썬)

다음 구현을 통해 더 잘 이해해 보겠습니다.

class Solution(object):
    def find(self, x):
        if self.parent[x] == -1:
            return x
        self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        parent_x = self.find(x)
        parent_y = self.find(y)
        if parent_x == parent_y:
            return
        self.parent[parent_y] = parent_x

    def minimumCost(self, n, connections):
        self.parent = [-1 for i in range(n + 1)]
        disjoint = n - 1
        cost = 0
        c = sorted(connections, key=lambda v: v[2])
        i = 0
        while i < len(c) and disjoint:
            x = c[i][0]
            y = c[i][1]
            if self.find(x) != self.find(y):
                disjoint -= 1
                cost += c[i][2]
                self.union(x, y)
            i += 1
        return cost if not disjoint else -1

ob = Solution()
print(ob.minimumCost(3, [[1, 2, 5], [1, 3, 6], [2, 3, 1]]))

입력

3
[[1,2,5],[1,3,6],[2,3,1]]

출력

동작 과정을 살펴보겠습니다. 간선을 비용 기준으로 정렬하면 [2, 3, 1] → [1, 2, 5] → [1, 3, 6] 순서가 됩니다. 먼저 비용이 1인 간선 (2, 3)을 선택하고(disjoint: 2 → 1, cost = 1), 다음으로 비용이 5인 간선 (1, 2)을 선택하면(disjoint: 1 → 0, cost = 6) 세 도시가 모두 연결됩니다. 마지막 간선 (1, 3)은 이미 같은 집합에 속하므로 건너뛰며, 최종 결과는 6입니다.

복잡도 분석

간선 정렬에 O(E log E)의 시간이 걸리고, 유니온-파인드 연산은 경로 압축 덕분에 거의 상수 시간에 처리되므로 전체 시간 복잡도는 O(E log E)입니다. 공간 복잡도는 parent 배열로 인해 O(N)입니다.