구현: DSU 클래스 만들기

이 장에서는 경로 압축과 크기 기준 합치기를 모두 갖춘 DSU 클래스를 Python으로 작성하고 한 줄씩 설명합니다. 이어서 같은 핵심 루틴을 C++, Java, TypeScript로 옮깁니다. 네 언어 모두 원소를 0부터 n-1까지의 정수로 다룹니다.

Python 구현

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

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

    def union(self, a: int, b: int) -> bool:
        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]
        self.count -= 1
        return True

    def connected(self, a: int, b: int) -> bool:
        return self.find(a) == self.find(b)

    def set_size(self, x: int) -> int:
        return self.size[self.find(x)]

한 줄씩 살펴보기

• self.parent = list(range(n)): 모든 원소가 자기 자신을 부모로 가리키므로 처음에는 n개의 집합이 있습니다.
• self.size = [1] * n: 각 집합의 원소 수입니다. 루트 위치의 값만 의미가 있습니다.
• self.count = n: 현재 집합 개수입니다. 합치기에 성공할 때마다 1씩 줄어듭니다.
• find의 첫 번째 반복문은 부모를 따라 올라가 루트를 찾습니다.
• 두 번째 반복문은 경로 압축입니다. 지나온 노드의 부모를 루트로 바꾸고 다음 노드로 이동합니다. 오른쪽 값이 먼저 계산되므로 self.parent[x], x = root, self.parent[x]는 원래 부모로 안전하게 이동합니다.
• 재귀 대신 반복문을 쓴 이유는 Python의 기본 재귀 한도(약 1000)에 걸리지 않기 위해서입니다.
• union은 먼저 두 루트를 찾고, 같으면 이미 같은 집합이므로 False를 돌려줍니다. 이 반환값이 사이클 판별에 그대로 쓰입니다.
• 크기가 작은 쪽을 rb로 맞춘 뒤 rb를 ra 밑에 붙이고 크기를 더합니다.
• connected와 set_size는 find를 감싼 편의 메서드입니다.

사용 예

dsu = DSU(7)
edges = [(0, 1), (1, 2), (3, 4), (5, 6), (2, 0)]
for a, b in edges:
    if not dsu.union(a, b):
        print("사이클을 만드는 간선:", (a, b))   # (2, 0)

print(dsu.count)            # 3
print(dsu.connected(0, 2))  # True
print(dsu.set_size(4))      # 2

C++

#include <numeric>
#include <utility>
#include <vector>

struct DSU {
    std::vector<int> parent, size;
    int count;

    explicit DSU(int n) : parent(n), size(n, 1), count(n) {
        std::iota(parent.begin(), parent.end(), 0);
    }

    int find(int x) {
        int root = x;
        while (parent[root] != root) root = parent[root];
        while (parent[x] != root) {
            int next = parent[x];
            parent[x] = root;
            x = next;
        }
        return root;
    }

    bool unite(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra == rb) return false;
        if (size[ra] < size[rb]) std::swap(ra, rb);
        parent[rb] = ra;
        size[ra] += size[rb];
        --count;
        return true;
    }
};

C++에서 union은 예약어이므로 메서드 이름을 unite로 했습니다.

Java

public final class DSU {
    private final int[] parent;
    private final int[] size;
    private int count;

    public DSU(int n) {
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            size[i] = 1;
        }
        count = n;
    }

    public int find(int x) {
        int root = x;
        while (parent[root] != root) root = parent[root];
        while (parent[x] != root) {
            int next = parent[x];
            parent[x] = root;
            x = next;
        }
        return root;
    }

    public boolean union(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra == rb) return false;
        if (size[ra] < size[rb]) { int t = ra; ra = rb; rb = t; }
        parent[rb] = ra;
        size[ra] += size[rb];
        count--;
        return true;
    }

    public int count() { return count; }
}

TypeScript

export class DSU {
  private parent: Int32Array;
  private size: Int32Array;
  count: number;

  constructor(n: number) {
    this.parent = Int32Array.from({ length: n }, (_, i) => i);
    this.size = new Int32Array(n).fill(1);
    this.count = n;
  }

  find(x: number): number {
    let root = x;
    while (this.parent[root] !== root) root = this.parent[root];
    while (this.parent[x] !== root) {
      const next = this.parent[x];
      this.parent[x] = root;
      x = next;
    }
    return root;
  }

  union(a: number, b: number): boolean {
    let ra = this.find(a);
    let rb = this.find(b);
    if (ra === rb) return false;
    if (this.size[ra] < this.size[rb]) [ra, rb] = [rb, ra];
    this.parent[rb] = ra;
    this.size[ra] += this.size[rb];
    this.count--;
    return true;
  }
}

TypeScript에서는 Int32Array를 써서 숫자 배열을 촘촘하게 저장합니다. 원소가 문자열이라면 Map<string, number>로 먼저 번호를 붙인 뒤 같은 클래스를 쓰면 됩니다.

정리

DSU는 parent와 size 배열, 그리고 find와 union 두 메서드만으로 완성됩니다. find는 반복문으로 루트를 찾고 경로를 압축하며, union은 작은 집합을 큰 집합 밑에 붙이고 성공 여부를 돌려줍니다. 네 언어의 구현 모두 구조가 같으므로 하나를 이해하면 나머지도 그대로 읽힙니다.