[백준] 1717번: 집합의 표현 (파이썬)

✅문제: 1717번

1717번: 집합의 표현

📌개념정리

(1) 분리 집합(Disjoint Set)

• 정의: 상호 배타적인 부분 집합들로 나누어진 전체 집합을 효율적으로 표현하고 처리하기 위한 자료구조
• 주요 연산
  • find: 특정 원소가 속한 집합을 찾는 연산
  • union: 두 원소가 속한 집합을 하나로 합치는 연산
• 구현 방법
  • 경로 압축(Path Compression): find 함수를 재귀적으로 호출한 뒤, 찾은 루트 노드가 바로 부모 노드가 되도록 함.
  • 합병(by rank 또는 size): 두 트리를 하나로 합칠 때, 더 작은 높이 또는 크기의 트리를 더 큰 트리의 서브트리로 만듦.

📌문제풀이

주어진 n개의 원소에 대해 두 종류의 연산을 수행하는 과정을 시뮬레이션하는 문제임. 연산 0은 두 원소를 같은 집합에 합치는 union 연산이고, 연산 1은 두 원소가 같은 집합에 속하는지를 확인하는 find 연산임.

1. 주어진 입력을 읽고 각 연산을 적절하게 분류하여 수행
2. 분리 집합 자료구조를 사용하여 각 원소의 집합 관계를 유지
3. union 연산은 두 원소의 루트 노드를 찾아 루트 노드를 하나로 설정
4. find 연산은 두 원소의 루트 노드가 같은지 확인하여 같다면 "YES", 다르다면 "NO"를 결과에 추가

import sys
input = sys.stdin.read

# 주어진 원소 x의 루트 노드를 찾아 경로 압축
def find(parent, x):
    if parent[x] != x:
        parent[x] = find(parent, parent[x])  # 경로 압축
    return parent[x]

# 두 원소 a, b를 같은 집합으로 합침
def union(parent, a, b):
    rootA = find(parent, a)
    rootB = find(parent, b)
    if rootA != rootB:
        parent[rootB] = rootA  # 두 트리 합침

def solve():
    input_data = input().split()
    index = 0
    n = int(input_data[index])  # 원소 개수
    index += 1
    m = int(input_data[index])  # 연산 개수
    index += 1
    
    parent = list(range(n + 1))  # 부모 테이블 초기화
    result = []

    for _ in range(m):
        operation = int(input_data[index])  # 연산
        index += 1
        a = int(input_data[index])  # 첫 번째 원소
        index += 1
        b = int(input_data[index])  # 두 번째 원소
        index += 1
        
        if operation == 0:  # union 연산
            union(parent, a, b)
        elif operation == 1:  # find 연산
            if find(parent, a) == find(parent, b):
                result.append("YES")
            else:
                result.append("NO")

    print("\n".join(result))

if __name__ == "__main__":
    solve()

코드 실행 결과를 예제 입력 1로 확인하면, 처음에 n과 m, m개의 연산이 주어짐.

# 예제 입력 1
7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

첫 줄에서 7은 원소의 개수이고, 8은 연산의 수임. 연산은 0(union)과 1(find) 두 가지가 있음.

 

모든 원소는 자기 자신을 부모로 가지는 집합임. 각각의 연산을 진행하면 다음와 같음.

원소	1	2	3	4	5	6	7
부모	1	2	3	4	5	6	7

1) 0 1 3 연산: union(1, 3)

원소	1	2	3	4	5	6	7
부모	1	2	1	4	5	6	7

2) 1 1 7 연산: find(1, 7)
1과 7이 같은 집합에 속하는지 확인함. 포함되어 있지 않으므로 "NO"를 출력함.

3) 0 7 6 연산: union(7, 6)

원소	1	2	3	4	5	6	7
부모	1	2	1	4	5	7	7

4) 1 7 1 연산: find(7, 1)
7과 1이 같은 집합에 속하는지 확인함. 포함되어 있지 않으므로 "NO"를 출력함.

5) 0 3 7 연산: union(3, 7)

원소	1	2	3	4	5	6	7
부모	1	2	1	4	5	7	1

6) 0 4 2 연산: union(4, 2)

원소	1	2	3	4	5	6	7
부모	1	4	1	4	5	7	1

7) 0 1 1 연산: union(1, 1)
같은 원소에 대한 요청이므로 변경사항 없음.

8) 1 1 1 연산: find(1, 1)
같은 원소에 대한 요청이므로 "YES"를 출력함.

# 예제 출력 1
NO
NO
YES

📌후기

• 분리 집합은 집합 간의 관계를 처리하는 자료구조라는 것을 학습했음.
• 문제를 통해 경로 압축과 합병 기법을 사용하는 법을 이해함.

📌참고자료

1) 서녁, "[Python] 분리집합, 서로소집합", 2022.04.26, https://velog.io/@ashooozzz/Python-분리집합-서로소집합

2) 더북, "파이썬으로 배우는 자료 구조 핵심 원리: 13.2.2 분리 집합", https://thebook.io/080200/0282/

3) 4Legs, "분리 집합 (Disjoint Set) : Union-Find", 2020.12.28, https://4legs-study.tistory.com/94