[백준] 1717번 집합의 표현 해설 및 풀이 (Python)

백준 1717번 집합의 표현

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

 

---

문제

초기에 n+1개의 집합 {0},{1},{2},…,{n}{0}, {1}, {2}, ... , {n}이 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

 

집합을 표현하는 프로그램을 작성하시오.

입력

첫째 줄에 n, m이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다.

출력

1로 시작하는 입력에 대해서 a와 b가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

제한

• 1 ≤ n ≤ 1,000,000 
• 1 ≤ m ≤ 100,000
• 0 ≤ a, b ≤ n
•  a, b는 정수
•  a와 b는 같을 수도 있다.

 

문제 유형

자료 구조

분리 집합

 

풀이 방법 도출

 

합집합 연산과 같은 집합인지 확인하는 Union-Find 을 사용하는 것이 이번 문제의 풀이 방법입니다.

 

1. 각 원소를 자기 자신을 부모로 설정하고 초기화합니다.
2. 재귀로 부모를 찾으면서, 찾은 부모를 현재 노드에 바로 연결합니다. 이후 트리의 높이 감소합니다. 트리의 높이가 감소하며 탐색 속도 향상됩니다.
3. 두 원소의 부모를 찾아서 한쪽 부모를 다른 쪽 부모에 연결합니다.
4. 0 a b 합집합이면 union(a, b)를, 1 a b 같은 집합인지 확인하고 find(a)와 find(b)가 같은지 비교 후 YES/NO를 출력합니다.

시간 복잡도

• Union-Find 연산 : O(N)
• 전체 시간 복잡도 : O(M * N)

---

핵심 코드 삽입 및 설명

import sys
sys.setrecursionlimit(10**7)  

input = sys.stdin.readline

# 부모를 찾는 함수
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])  # 경로 압축
    return parent[x]

# 두 집합을 합치는 함수
def union(x, y):
    px = find(x)
    py = find(y)
    if px != py:
        parent[py] = px  # 한 쪽 부모로 합침

# 원소 개수 n, 연산 개수 m 입력
n, m = map(int, input().split())

# 초기: 자기 자신이 부모
parent = [i for i in range(n + 1)]

for _ in range(m):
    op, a, b = map(int, input().split())
    if op == 0:
        # 합집합 연산
        union(a, b)
    else:
        # 같은 집합인지 확인
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")