유니온 파인드(Union-Find) - Python

Why?

왜 쓰냐 => 여러 개의 원소가 있을 때 다음 두 가지 작업을 효율적으로 수행하기 위해 사용

 

Union? Find?

1. Union : 두 개의 집합을 합침

2. Find : 특정 원소가 속한 집합을 찾음

 

주로 그래프 알고리즘, 사이클 검출, 최소 신장 트리 등의 문제에 사용

 

 

---

핵심 개념

 

1. 부모 노드(parent): 각 원소는 자신이 속한 집합을 대표하는 "부모 노드"를 가집니다.

 

• 초기에는 모든 원소가 자기 자신을 부모로 가집니다.
• Find 연산을 통해 부모를 찾을 수 있습니다.

 

 

2. 경로 압축(Path Compression): Find 연산을 수행하면서 트리의 깊이를 줄여 효율성을 높이는 기법입니다.

• 경로 압축을 통해 원소들이 직접 루트 노드를 가리키게 되어 Find 연산이 매우 빠르게 수행됩니다.

 

3. 랭크(rank) 또는 크기(size): Union 연산에서 두 집합을 합칠 때 트리의 깊이를 최소화하기 위해 랭크 또는 크기를 기준으로 합칩니다.

 

 

4. 시간 복잡도 : 거의 O(1)

 

---

구현

 

class UnionFind:
    def __init__(self, n):
        # 부모 노드를 자기 자신으로 초기화
        self.parent = [i for i in range(n)]
        # 각 집합의 랭크(깊이)를 저장
        self.rank = [0] * n

    def find(self, x):
        # 경로 압축 적용
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        # 두 원소의 루트 노드를 찾음
        root_x = self.find(x)
        root_y = self.find(y)

        # 이미 같은 집합에 속해있으면 합치지 않음
        if root_x == root_y:
            return

        # 랭크(깊이)가 낮은 쪽을 높은 쪽에 붙임
        if self.rank[root_x] < self.rank[root_y]:
            self.parent[root_x] = root_y
        elif self.rank[root_x] > self.rank[root_y]:
            self.parent[root_y] = root_x
        else:
            # 랭크가 같으면 임의로 하나를 다른 쪽에 붙이고 랭크를 증가시킴
            self.parent[root_y] = root_x
            self.rank[root_x] += 1

    def connected(self, x, y):
        # 두 원소가 같은 집합에 속해있는지 확인
        return self.find(x) == self.find(y)


# 테스트 코드
if __name__ == "__main__":
    uf = UnionFind(10)  # 원소 0부터 9까지 10개

    # 일부 집합 병합
    uf.union(1, 2)
    uf.union(2, 3)
    uf.union(4, 5)
    uf.union(6, 7)
    uf.union(5, 6)

    # 연결 여부 확인
    print(uf.connected(1, 3))  # True (같은 집합)
    print(uf.connected(1, 4))  # False (다른 집합)
    print(uf.connected(4, 7))  # True (같은 집합)

    # 부모 확인
    print(uf.find(3))  # 1 또는 2가 대표 노드
    print(uf.find(7))  # 4 또는 5가 대표 노드

 

# 초기화 함수
def initialize(n):
    global parent, rank
    parent = [i for i in range(n)]  # 부모 노드를 자기 자신으로 초기화
    rank = [0] * n  # 각 집합의 랭크(깊이)를 0으로 초기화

# Find 함수 (경로 압축 적용)
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])  # 경로 압축
    return parent[x]

# Union 함수 (랭크 기반 병합)
def union(x, y):
    root_x = find(x)
    root_y = find(y)

    if root_x == root_y:  # 이미 같은 집합에 속해 있다면 합치지 않음
        return

    # 랭크(깊이)가 낮은 쪽을 높은 쪽에 붙임
    if rank[root_x] < rank[root_y]:
        parent[root_x] = root_y
    elif rank[root_x] > rank[root_y]:
        parent[root_y] = root_x
    else:
        parent[root_y] = root_x
        rank[root_x] += 1  # 랭크가 같으면 한 쪽의 랭크를 증가시킴

# 연결 여부 확인 함수
def connected(x, y):
    return find(x) == find(y)

# 테스트 코드
if __name__ == "__main__":
    n = 10  # 원소 개수
    initialize(n)  # 유니온 파인드 초기화

    # 일부 집합 병합
    union(1, 2)
    union(2, 3)
    union(4, 5)
    union(6, 7)
    union(5, 6)

    # 연결 여부 확인
    print(connected(1, 3))  # True (1, 2, 3이 같은 집합)
    print(connected(1, 4))  # False (1과 4는 다른 집합)
    print(connected(4, 7))  # True (4, 5, 6, 7이 같은 집합)

    # 부모 확인
    print(find(3))  # 1 또는 2가 대표 노드
    print(find(7))  # 4 또는 5가 대표 노드