Python으로 마을 물 공급 비용 최적화하기: MST와 Union-Find 활용법

문제 소개

마을에 n개의 집이 있다고 가정해 보겠습니다. 우리는 우물을 짓고 파이프를 깔아 모든 집에 물을 공급해야 하며, 각 집 i에는 두 가지 선택지가 있습니다.

집 안에 우물 짓기: 비용은 wells[i]로 주어집니다.
다른 우물에서 파이프로 물 끌어오기: 집 사이 파이프 설치 비용은 배열 pipes로 주어지며, 각 pipes[i]는 [house1, house2, cost] 형태로 house1과 house2를 연결하는 비용을 의미합니다. 연결은 양방향입니다.

목표는 모든 집에 물을 공급하는 최소 총비용을 구하는 것입니다. 예를 들어 입력이 n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]이라면 정답은 3입니다.

위 그림처럼 첫 번째 집에 비용 1로 우물을 짓고, 나머지 두 집을 각각 비용 1짜리 파이프로 연결하면 총비용 3으로 모든 집에 물을 공급할 수 있습니다. 참고로 이 문제는 LeetCode 1168번 'Optimize Water Distribution in a Village'와 동일한 문제입니다.

접근 방법: 가상 노드 + 크루스칼 알고리즘

이 문제는 최소 신장 트리(MST)로 변환해 풀 수 있는 대표적인 유형입니다. 핵심 트릭은 가상 노드 0을 도입하는 것입니다. 각 집 i를 노드 0과 wells[i] 비용의 간선으로 연결하면 '우물 짓기'도 하나의 간선 선택으로 취급할 수 있습니다. 이후 모든 간선(파이프 + 우물)을 비용순으로 정렬하고, 크루스칼 알고리즘으로 사이클을 만들지 않는 간선만 선택하면 최소 비용을 구할 수 있습니다.

알고리즘 단계

find() 함수를 정의합니다. 이 함수는 노드 a의 루트 부모를 찾습니다.
parent[a]가 -1이면 a 자체가 루트이므로 a를 반환합니다.
그렇지 않으면 parent[a] := find(parent[a])로 경로 압축을 수행한 뒤 parent[a]를 반환합니다.
union() 함수를 정의합니다. 이 함수는 두 노드 a, b를 하나의 집합으로 합칩니다.
parent_a := find(a), parent_b := find(b)
parent_a == parent_b라면 이미 같은 집합이므로 True를 반환합니다(사이클 발생).
그렇지 않으면 parent[parent_b] := parent_a로 병합하고 False를 반환합니다.
메인 로직은 다음과 같습니다.
크기가 n + 1인 리스트 parent를 만들고 모든 값을 -1로 초기화합니다.
각 우물 비용 wells[i]에 대해 간선 [0, i+1, wells[i]]를 pipes에 추가합니다. 여기서 노드 0은 가상 우물 노드입니다.
pipes 배열을 비용 기준으로 오름차순 정렬합니다.
cost := 0으로 초기화한 뒤, pipes의 각 간선에 대해 union(source, destination) 결과가 False일 때(실제로 새로 연결될 때)만 cost에 해당 비용을 더합니다.
최종 cost를 반환합니다.

Python 구현 예제

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

    def union(self, a, b):
        parent_a = self.find(a)
        parent_b = self.find(b)
        if parent_a == parent_b:
            return True
        self.parent[parent_b] = parent_a
        return False

    def minCostToSupplyWater(self, n, well, pipes):
        self.parent = [-1 for i in range(n + 1)]
        for i in range(len(well)):
            pipes.append([0, i + 1, well[i]])
        pipes = sorted(pipes, key=lambda v: v[2])
        cost = 0
        for edge in pipes:
            source = edge[0]
            destination = edge[1]
            temp = edge[2]
            if not self.union(source, destination):
                cost += temp
        return cost

ob = Solution()
print(ob.minCostToSupplyWater(3, [1, 2, 2], [[1, 2, 1], [2, 3, 1]]))

입력

3, [1,2,2], [[1,2,1],[2,3,1]]

출력

복잡도 분석

간선의 개수를 E라 하면, 정렬에 O(E log E)의 시간이 소요되고 유니온-파인드 연산은 경로 압축 덕분에 거의 상수 시간에 처리되므로 전체 시간 복잡도는 O(E log E)입니다. 공간 복잡도는 parent 배열 저장을 위해 O(V)입니다.