[백준 1717번 유니온 파인드 알고리즘] 집합의 표현 파이썬 풀이

문제 풀이하기전에 Union-find에 대해 살펴보겠습니다.

유니온 파인드란 간단하게

두 노드, 두 숫자, 두 무언가가 같은 집합안에 있는지 확인하는 것입니다.

 

1

2

3

4

5

6

7

8

9

10

 

여기에 원본 배열이 있습니다.

이때

1 - 3 - 6이 연결되어있다고 가정을 하면

값이 가르키고 있는 부모배열은 이렇게 표현할 수 있습니다.

1

2

1

4

5

3

7

8

9

10

 

6은 3 , 3은 1을 가르키고 1은 1자신을 가르키고 있죠?

이때1,3,6은 같은 집합으로 연결되어있다고 표현할 수 있습니다.

 

즉 유니온 - 파인드는 

유니온은 저 배열을 만들어주는것이고

파인드는 저 배열속에서 찾는것이라고 생각하면 간편합니다. 

 

import sys
input = sys.stdin.readline

def _union(A,B):
    par[B] = A

def _find(A):
    if par[A] == A:
        return A
    else:
        return _find(par[A])

N,M = map(int, input().strip().split())
par = [i for i in range(N+1)]

for _ in range(M):
    X,A,B = map(int, input().strip().split())

    if X == 0:
        _union(A,B)
    else:
        if _find(A) == _find(B):
            print("YES")
        else:
            print("NO")

 

이것이 유니온파인드의 기초 코드 및 로직입니다.

이 코드를 먼저 유심히 보면서 이해하고 다음 코드를 봐주세요.

 

왜 최적화를 해야할까요?

만약 트리가 일렬로 나란히 1억개가 있다고하면

그 1억개를 전부 재귀함수를 돌려서 찾아내야해요.

너무나 비효율적일것같지 않나요?

 

import sys
input = sys.stdin.readline

def _union(A,B):
    A = _find(A)
    B = _find(B)

    if A == B:
        return
    
    if rank[A] < rank[B]:
        par[A] = B
    elif rank[B] < rank[A]:
        par[B] = A
    else:
        par[B] = A
        rank[A] += 1


def _find(A):
    if par[A] != A:
        par[A] = _find(par[A])  # 경로 압축
    return par[A]

N,M = map(int, input().strip().split())
rank = [0 for _ in range(N+1)]
par = [i for i in range(N+1)]

for _ in range(M):
    X,A,B = map(int, input().strip().split())

    if X == 0:
        _union(A,B)
    else:
        if _find(A) == _find(B):
            print("YES")
        else:
            print("NO")

 

코드 설명

1. rank 배열
   • rank[i]는 i가 루트인 트리의 깊이를 나타냅니다. 이는 트리의 높이를 최소화하여 유니온-파인드 연산의 효율성을 높이기 위해 사용됩니다.
   • rank 배열을 이용하면 두 트리를 합칠 때, 더 깊은 트리 밑에 얕은 트리를 붙임으로써 트리의 높이가 불필요하게 커지는 것을 방지할 수 있습니다.
2. par 배열
   • par[i]는 i의 부모 노드를 나타냅니다. 초기에는 각 원소가 자기 자신을 부모로 가지도록 설정됩니다.
3. _union(A, B) 함수
   • A와 B가 속한 두 집합을 하나로 합칩니다.
   • 먼저 A와 B의 루트 노드를 찾습니다.
   • 두 루트 노드의 rank를 비교하여, 더 작은 rank를 가진 트리를 더 큰 rank를 가진 트리 밑에 붙입니다.
   • 만약 두 트리의 rank가 같다면, 임의로 한 트리를 다른 트리 밑에 붙이고 그 트리의 rank를 1 증가시킵니다.
4. _find(A) 함수
   • A가 속한 집합의 루트 노드를 찾습니다.
   • 경로 압축(path compression)을 적용하여, 찾기 연산을 할 때 경로 상의 모든 노드를 직접 루트에 연결시킵니다. 이로써 다음 찾기 연산이 더 빠르게 수행됩니다.

예시

• 초기 상태
  • rank: [0, 0, 0, 0, ..., 0] (모든 노드의 랭크는 0)
  • par: [0, 1, 2, 3, ..., N] (모든 노드는 자기 자신을 부모로 가짐)
• 연산 예시
  1. Union 연산: _union(1, 2)
     • 1과 2의 루트를 각각 찾습니다. (둘 다 자기 자신이 루트)
     • rank가 동일하므로 2를 1의 자식으로 하고, 1의 rank를 1 증가시킵니다.
     • rank: [0, 1, 0, 0, ..., 0]
     • par: [0, 1, 1, 3, ..., N]
  2. Find 연산: _find(2)
     • 2의 부모는 1이므로, 1을 루트로 반환합니다.
     • 경로 압축에 의해 2의 부모를 1로 설정합니다.
  3. Union 연산: _union(2, 3)
     • 2와 3의 루트를 찾습니다. (2의 루트는 1, 3의 루트는 3)
     • rank가 작은 3을 1의 자식으로 설정합니다.
     • rank는 변하지 않습니다.
     • par: [0, 1, 1, 1, ..., N]

 

이 코드 그냥 외우지 마시고 꼭 손으로 한번 써보면서 이해해보시기를 추천드립니다!