연습 문제

이 장의 문제 네 개는 모두 유니온 파인드의 대표적인 쓰임새를 하나씩 담고 있습니다. 각 문제에 접근 방법과 Python 풀이를 붙였습니다. 풀이는 아래의 짧은 DSU를 공통으로 사용합니다.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True

문제 1. 동아리 수 세기

학생 n명(0번부터 n-1번)과 "두 학생이 같은 동아리"라는 쌍 목록이 주어집니다. 같은 동아리 관계는 전이적입니다. 동아리가 모두 몇 개인지, 가장 큰 동아리의 인원은 몇 명인지 구하세요.

접근: 쌍마다 union을 호출하고, 마지막에 서로 다른 루트의 개수와 루트 위치의 size 최댓값을 구합니다.

def clubs(n, pairs):
    dsu = DSU(n)
    for a, b in pairs:
        dsu.union(a, b)
    roots = {dsu.find(x) for x in range(n)}
    return len(roots), max(dsu.size[r] for r in roots)

print(clubs(6, [(0, 1), (1, 2), (4, 5)]))  # (3, 3)

문제 2. 쓸모없는 케이블 찾기

사무실 컴퓨터 n대를 케이블로 잇는 작업 기록이 순서대로 주어집니다. 이미 연결된 두 컴퓨터를 다시 잇는 케이블은 고리를 만들 뿐 쓸모가 없습니다. 쓸모없는 케이블의 번호를 기록 순서대로 모두 출력하세요.

접근: 무방향 그래프의 사이클 판별입니다. union이 False를 돌려주면 두 끝점이 이미 같은 집합이므로 그 간선이 사이클을 만듭니다.

def useless_cables(n, cables):
    dsu = DSU(n)
    return [i for i, (a, b) in enumerate(cables) if not dsu.union(a, b)]

print(useless_cables(4, [(0, 1), (1, 2), (2, 0), (2, 3), (3, 1)]))  # [2, 4]

문제 3. 계정 합치기

각 계정은 이름과 이메일 목록으로 이루어집니다. 이메일을 하나라도 공유하는 계정은 같은 사람의 것이며, 이 관계도 전이적입니다. 같은 사람의 계정을 합쳐 사람마다 정렬된 이메일 목록을 만드세요.

접근: 계정 번호를 원소로 삼습니다. 이메일마다 처음 본 계정을 기억해 두고, 같은 이메일이 다시 나오면 두 계정을 합칩니다. 마지막에 루트별로 이메일을 모읍니다.

from collections import defaultdict

def merge_accounts(accounts):
    dsu = DSU(len(accounts))
    owner = {}
    for i, (_, emails) in enumerate(accounts):
        for e in emails:
            if e in owner:
                dsu.union(i, owner[e])
            else:
                owner[e] = i
    groups = defaultdict(set)
    for e, i in owner.items():
        groups[dsu.find(i)].add(e)
    return [(accounts[r][0], sorted(es)) for r, es in groups.items()]

accounts = [
    ("kim", ["a@x.io", "b@x.io"]),
    ("kim", ["c@x.io"]),
    ("kim", ["b@x.io", "d@x.io"]),
]
print(merge_accounts(accounts))
# [('kim', ['a@x.io', 'b@x.io', 'd@x.io']), ('kim', ['c@x.io'])]

문제 4. 가장 싼 도로망

도시 n개와 "도시 u와 v를 잇는 도로 건설 비용 w" 후보 목록이 주어집니다. 모든 도시를 연결하는 최소 비용을 구하고, 불가능하면 -1을 출력하세요.

접근: 크루스칼 알고리즘입니다. 비용이 싼 간선부터 보면서 사이클을 만들지 않는 간선만 고르면 최소 신장 트리가 됩니다. 간선 n-1개를 고르면 멈춥니다.

def cheapest_network(n, roads):
    dsu = DSU(n)
    total = used = 0
    for w, u, v in sorted((w, u, v) for u, v, w in roads):
        if dsu.union(u, v):
            total += w
            used += 1
            if used == n - 1:
                return total
    return total if n <= 1 else -1

roads = [(0, 1, 4), (0, 2, 1), (1, 2, 2), (2, 3, 5), (1, 3, 7)]
print(cheapest_network(4, roads))  # 8

정리

동아리 수 세기는 연결 요소, 쓸모없는 케이블은 사이클 판별, 계정 합치기는 키를 통한 그룹화, 도로망은 크루스칼 최소 신장 트리입니다. 네 문제 모두 union의 반환값과 루트별 집계만으로 풀립니다. 문제를 보면 "전이적으로 같은 묶음인가?"를 먼저 물어보세요.