[백준] / [Python] / [1717] 집합의 표현

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

---

문제

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

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

입력

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

출력

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

제한

•  1≤𝑛≤1000000
•  1≤𝑚≤100000
•  0≤𝑎,𝑏≤𝑛
•  𝑎, 𝑏는 정수
•  𝑎와 𝑏는 같을 수도 있다.

---

시간 초과

연결되어 있는 집합을 연산하면서 답을 구했다.

하지만 집합을 연산하는 과정에서 시간 복잡도가 조금 더 증가하였고 시간 초과를 출력하였다.

import sys
input = sys.stdin.readline

n, m = map(int, input().split())

set_num = [None] * (n+1)
set_list = []

for _ in range(m):
    x, a, b = map(int, input().split())

    if set_num[a] == None:
        set_list.append({a})
        set_num[a] = len(set_list) - 1
        a_idx = set_num[a]
    else:
        a_idx = set_num[a]
    
    if set_num[b] == None:
        set_list.append({b})
        set_num[b] = len(set_list) - 1
        b_idx = set_num[b]
    else:
        b_idx = set_num[b]
    
    a_idx, b_idx = min(a_idx, b_idx), max(a_idx, b_idx)
    if x == 0:
        set_list[a_idx] = set_list[a_idx] | set_list[b_idx]
        
        if a_idx != b_idx:
            for num in set_list[b_idx]:
                set_num[num] = a_idx

            set_list[b_idx] = None
        
    elif x == 1:
        if set_list[a_idx] == set_list[b_idx]:
            print("YES")
        else:
            print("NO")

정답

집합을 직접 구할 필요 없이 Union-Find 알고리즘으로 상수 시간에 문제를 해결할 수 있습니다.

import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

n, m = map(int, input().split())

parent = [i for i in range(n+1)]

def find_parent(x):
    if parent[x] == x:
        return parent[x]
    parent[x] = find_parent(parent[x])
    return parent[x]
    

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)
    
    if a < b: parent[b] = a
    else: parent[a] = b

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