15_mst.py - Examples

  1"""
  2유니온 파인드 (Union-Find / Disjoint Set Union)
  3Union-Find / Disjoint Set Union Data Structure
  4
  5서로소 집합을 관리하는 자료구조로, 그래프의 연결성 문제에 주로 사용됩니다.
  6"""
  7
  8from typing import List, Tuple
  9
 10
 11# =============================================================================
 12# 1. 기본 유니온 파인드
 13# =============================================================================
 14class UnionFind:
 15    """
 16    기본 유니온 파인드 구현
 17    - 경로 압축 (Path Compression)
 18    - 랭크 기반 합치기 (Union by Rank)
 19    """
 20
 21    def __init__(self, n: int):
 22        """
 23        n개의 요소로 초기화 (0 ~ n-1)
 24        """
 25        self.parent = list(range(n))  # 자기 자신을 부모로
 26        self.rank = [0] * n           # 트리의 높이 (근사값)
 27        self.count = n                # 집합의 개수
 28
 29    def find(self, x: int) -> int:
 30        """
 31        x가 속한 집합의 대표(루트) 찾기
 32        경로 압축으로 거의 O(1)
 33        """
 34        if self.parent[x] != x:
 35            self.parent[x] = self.find(self.parent[x])  # 경로 압축
 36        return self.parent[x]
 37
 38    def union(self, x: int, y: int) -> bool:
 39        """
 40        x와 y가 속한 집합을 합치기
 41        이미 같은 집합이면 False 반환
 42        """
 43        root_x = self.find(x)
 44        root_y = self.find(y)
 45
 46        if root_x == root_y:
 47            return False  # 이미 같은 집합
 48
 49        # 랭크 기반 합치기 (작은 트리를 큰 트리에 붙임)
 50        if self.rank[root_x] < self.rank[root_y]:
 51            self.parent[root_x] = root_y
 52        elif self.rank[root_x] > self.rank[root_y]:
 53            self.parent[root_y] = root_x
 54        else:
 55            self.parent[root_y] = root_x
 56            self.rank[root_x] += 1
 57
 58        self.count -= 1
 59        return True
 60
 61    def connected(self, x: int, y: int) -> bool:
 62        """x와 y가 같은 집합에 있는지 확인"""
 63        return self.find(x) == self.find(y)
 64
 65    def get_count(self) -> int:
 66        """현재 집합의 개수"""
 67        return self.count
 68
 69
 70# =============================================================================
 71# 2. 크기 기반 유니온 파인드
 72# =============================================================================
 73class UnionFindWithSize:
 74    """
 75    집합의 크기를 추적하는 유니온 파인드
 76    """
 77
 78    def __init__(self, n: int):
 79        self.parent = list(range(n))
 80        self.size = [1] * n  # 각 집합의 크기
 81
 82    def find(self, x: int) -> int:
 83        if self.parent[x] != x:
 84            self.parent[x] = self.find(self.parent[x])
 85        return self.parent[x]
 86
 87    def union(self, x: int, y: int) -> bool:
 88        root_x = self.find(x)
 89        root_y = self.find(y)
 90
 91        if root_x == root_y:
 92            return False
 93
 94        # 크기 기반 합치기 (작은 집합을 큰 집합에)
 95        if self.size[root_x] < self.size[root_y]:
 96            self.parent[root_x] = root_y
 97            self.size[root_y] += self.size[root_x]
 98        else:
 99            self.parent[root_y] = root_x
100            self.size[root_x] += self.size[root_y]
101
102        return True
103
104    def get_size(self, x: int) -> int:
105        """x가 속한 집합의 크기"""
106        return self.size[self.find(x)]
107
108
109# =============================================================================
110# 3. 연결 요소 개수
111# =============================================================================
112def count_components(n: int, edges: List[List[int]]) -> int:
113    """
114    n개의 노드와 간선 목록이 주어질 때 연결 요소 개수
115    """
116    uf = UnionFind(n)
117    for u, v in edges:
118        uf.union(u, v)
119    return uf.get_count()
120
121
122# =============================================================================
123# 4. 그래프에서 사이클 검출
124# =============================================================================
125def has_cycle(n: int, edges: List[List[int]]) -> bool:
126    """
127    무방향 그래프에서 사이클 존재 여부
128    간선을 추가할 때 이미 같은 집합이면 사이클
129    """
130    uf = UnionFind(n)
131    for u, v in edges:
132        if not uf.union(u, v):
133            return True  # 이미 연결됨 = 사이클
134    return False
135
136
137# =============================================================================
138# 5. 크루스칼 MST (최소 신장 트리)
139# =============================================================================
140def kruskal_mst(n: int, edges: List[Tuple[int, int, int]]) -> Tuple[int, List[Tuple[int, int, int]]]:
141    """
142    크루스칼 알고리즘으로 최소 신장 트리 구하기
143    edges: [(u, v, weight), ...]
144    반환: (총 가중치, MST 간선 리스트)
145    """
146    # 가중치 기준 정렬
147    edges = sorted(edges, key=lambda x: x[2])
148
149    uf = UnionFind(n)
150    mst_weight = 0
151    mst_edges = []
152
153    for u, v, w in edges:
154        if uf.union(u, v):
155            mst_weight += w
156            mst_edges.append((u, v, w))
157
158            # n-1개의 간선을 선택하면 완료
159            if len(mst_edges) == n - 1:
160                break
161
162    return mst_weight, mst_edges
163
164
165# =============================================================================
166# 6. 친구 관계 (계정 병합)
167# =============================================================================
168def merge_accounts(accounts: List[List[str]]) -> List[List[str]]:
169    """
170    같은 이메일을 가진 계정 병합
171    accounts[i] = [이름, 이메일1, 이메일2, ...]
172    """
173    from collections import defaultdict
174
175    # 이메일 -> 계정 인덱스 매핑
176    email_to_id = {}
177    email_to_name = {}
178
179    for i, account in enumerate(accounts):
180        name = account[0]
181        for email in account[1:]:
182            if email in email_to_id:
183                pass  # 나중에 union
184            email_to_id[email] = i
185            email_to_name[email] = name
186
187    # 유니온 파인드로 같은 사람의 계정 연결
188    n = len(accounts)
189    uf = UnionFind(n)
190
191    email_first_account = {}
192    for i, account in enumerate(accounts):
193        for email in account[1:]:
194            if email in email_first_account:
195                uf.union(i, email_first_account[email])
196            else:
197                email_first_account[email] = i
198
199    # 결과 집계
200    root_to_emails = defaultdict(set)
201    for i, account in enumerate(accounts):
202        root = uf.find(i)
203        for email in account[1:]:
204            root_to_emails[root].add(email)
205
206    # 결과 포맷팅
207    result = []
208    for root, emails in root_to_emails.items():
209        name = accounts[root][0]
210        result.append([name] + sorted(emails))
211
212    return result
213
214
215# =============================================================================
216# 7. 섬 연결하기 (2D 그리드)
217# =============================================================================
218def num_islands_union_find(grid: List[List[str]]) -> int:
219    """
220    '1'은 땅, '0'은 물
221    연결된 땅 덩어리(섬)의 개수
222    """
223    if not grid or not grid[0]:
224        return 0
225
226    rows, cols = len(grid), len(grid[0])
227
228    # 2D -> 1D 좌표 변환
229    def get_index(r, c):
230        return r * cols + c
231
232    uf = UnionFind(rows * cols)
233    land_count = 0
234
235    for r in range(rows):
236        for c in range(cols):
237            if grid[r][c] == '1':
238                land_count += 1
239                # 오른쪽, 아래 방향만 확인 (중복 방지)
240                for dr, dc in [(0, 1), (1, 0)]:
241                    nr, nc = r + dr, c + dc
242                    if 0 <= nr < rows and 0 <= nc < cols and grid[nr][nc] == '1':
243                        if uf.union(get_index(r, c), get_index(nr, nc)):
244                            land_count -= 1
245
246    return land_count
247
248
249# =============================================================================
250# 테스트
251# =============================================================================
252def main():
253    print("=" * 60)
254    print("유니온 파인드 (Union-Find) 예제")
255    print("=" * 60)
256
257    # 1. 기본 사용
258    print("\n[1] 기본 유니온 파인드")
259    uf = UnionFind(10)
260    operations = [(0, 1), (2, 3), (4, 5), (1, 2), (6, 7), (8, 9), (0, 9)]
261    for u, v in operations:
262        uf.union(u, v)
263        print(f"    union({u}, {v}) -> 집합 수: {uf.get_count()}")
264
265    print(f"\n    0과 9 연결됨? {uf.connected(0, 9)}")
266    print(f"    0과 6 연결됨? {uf.connected(0, 6)}")
267
268    # 2. 크기 추적
269    print("\n[2] 크기 기반 유니온 파인드")
270    uf_size = UnionFindWithSize(5)
271    uf_size.union(0, 1)
272    uf_size.union(2, 3)
273    uf_size.union(0, 2)
274    print(f"    0이 속한 집합 크기: {uf_size.get_size(0)}")
275    print(f"    4가 속한 집합 크기: {uf_size.get_size(4)}")
276
277    # 3. 연결 요소 개수
278    print("\n[3] 연결 요소 개수")
279    edges = [[0, 1], [1, 2], [3, 4]]
280    count = count_components(5, edges)
281    print(f"    노드 5개, 간선: {edges}")
282    print(f"    연결 요소 개수: {count}")
283
284    # 4. 사이클 검출
285    print("\n[4] 사이클 검출")
286    edges_no_cycle = [[0, 1], [1, 2], [2, 3]]
287    edges_with_cycle = [[0, 1], [1, 2], [2, 0]]
288    print(f"    간선 {edges_no_cycle}: 사이클 = {has_cycle(4, edges_no_cycle)}")
289    print(f"    간선 {edges_with_cycle}: 사이클 = {has_cycle(3, edges_with_cycle)}")
290
291    # 5. 크루스칼 MST
292    print("\n[5] 크루스칼 MST")
293    #     1
294    #   0---1
295    #   |\  |
296    # 4 | \ |2
297    #   |  \|
298    #   3---2
299    #     3
300    edges_mst = [
301        (0, 1, 1), (0, 2, 4), (0, 3, 4),
302        (1, 2, 2), (2, 3, 3)
303    ]
304    total_weight, mst_edges = kruskal_mst(4, edges_mst)
305    print(f"    간선: {edges_mst}")
306    print(f"    MST 총 가중치: {total_weight}")
307    print(f"    MST 간선: {mst_edges}")
308
309    # 6. 계정 병합
310    print("\n[6] 계정 병합")
311    accounts = [
312        ["John", "john@mail.com", "john_work@mail.com"],
313        ["John", "john@mail.com", "john2@mail.com"],
314        ["Mary", "mary@mail.com"],
315        ["John", "john3@mail.com"]
316    ]
317    result = merge_accounts(accounts)
318    print(f"    입력:")
319    for acc in accounts:
320        print(f"      {acc}")
321    print(f"    병합 결과:")
322    for acc in result:
323        print(f"      {acc}")
324
325    # 7. 섬 개수 (유니온 파인드)
326    print("\n[7] 섬의 개수 (유니온 파인드)")
327    grid = [
328        ['1', '1', '0', '0', '0'],
329        ['1', '1', '0', '0', '0'],
330        ['0', '0', '1', '0', '0'],
331        ['0', '0', '0', '1', '1']
332    ]
333    count = num_islands_union_find(grid)
334    print(f"    격자:")
335    for row in grid:
336        print(f"    {row}")
337    print(f"    섬의 개수: {count}")
338
339    print("\n" + "=" * 60)
340    print("유니온 파인드 시간 복잡도")
341    print("=" * 60)
342    print("""
343    경로 압축 + 랭크/크기 기반 합치기 사용 시:
344    - find(): 거의 O(1) (정확히는 O(α(n)), α는 아커만 역함수)
345    - union(): 거의 O(1)
346    - 공간 복잡도: O(n)
347
348    주요 활용:
349    - 연결 요소 관리
350    - 사이클 검출
351    - 최소 신장 트리 (크루스칼)
352    - 동적 연결성 문제
353    """)
354
355
356if __name__ == "__main__":
357    main()