[백준] 1717번: 집합의 표현 - Python

링크: https://www.acmicpc.net/problem/1717

 

아이디어

• 초기에 n+1개의 집합이 존재하고, 합집합과 동일 집합 여부 판별 연산을 빠르게 처리해야 한다.
• 입력의 개수가 최대 10만 개이고, 노드 수가 최대 100만 개이므로 DFS나 재귀는시간초과가 발생한다.
• 집합 관리 알고리즘인 유니온 파인드 자료구조를 사용하자.
• 경로 압축 기법을 함께 사용하면 거의 상수 시간 복잡도로 처리할 수 있도록 하자.
• 배열을 만들어 각 원소의 부모 노드를 저장한다.
• find(x) 함수로 x의 루트를 찾고 경로 압축을 수행해 트리 깊이를 줄인다.
• union(a, b) 함수는 두 집합을 하나로 합친다.
• find(a) == find(b)이면 같은 집합에 속해 있다는 의미로 YES 출력

구현

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
li = [i for i in range(n + 1)]

def find(x):
    if li[x] != x:
        li[x] = find(li[x])  # 경로 압축
    return li[x]

def union(a, b):
    a_root = find(a)
    b_root = find(b)
    if a_root != b_root:
        li[b_root] = a_root

for _ in range(m):
    f, a, b = map(int, input().split())
    if f == 0:
        union(a, b)
    else:
        print("YES" if find(a) == find(b) else "NO")

마무리

어우 난생 처음 접해보는 알고리즘이었다. 처음엔 단순히 BFS로 풀 수 있을 줄 알고 접근했는데, 아니나 다를까 시간 초과가 났다.
서칭을 해보니 이건 유니온 파인드 알고리즘을 써야 하는 문제였다. 처음 보는 구조라 어렵게 느껴졌지만, 핵심 개념은 의외로 단순했다.
각 노드가 어떤 집합에 속해 있는지를 트리 구조로 표현하고,
합집합 연산으로 두 트리를 연결
같은 집합인지 확인 연산을 통해 루트를 비교.

여기에 경로 압축을 적용하면 탐색하면서 만나는 모든 노드들의 부모를 루트로 바로잡기 때문에 이후 연산 속도도 확연히 빨라진다.
이게 유니온 파인드의 진짜 핵심 포인트같다.

정말 자주 쓰이는 알고리즘이라고 하던데 왜 그런지 이번 문제를 통해 확실히 느꼈다.

쉽게 이해할 수 있도록 유니온 파인드를 시각화한 글도 따로 올려놨다.

참고하면 분명 도움될 거다.

https://yeachan.tistory.com/159

[알고리즘] Union-Find find() 함수 시각화