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

문제

0에서 n까지 숫자 한개로 이루어진 n+1개의 집합이 있다.

합집합 연산(0)과 두 원소가 같은 집합에 포함되는지 여부를 확인하는 연산(1)을 수행할 때,

같은 집합에 포함되면 yes, 아니면 no를 출력하는 문제이다.

연산은 연산종류(0,1), a원소, b원소 순서대로 입력된다.

풀이

• 이 문제는 각 원소가 포함된 집합의 대표 번호(부모 노드)를 저장함으로써
• 합집합 연산에서는 두 원소의 부모 노드를 더 작은 번호로 일치시켜주고 (union)
• 두 원소의 집합이 같은지 확인하는 연산에서는 부모 노드가 동일한지 확인하면 되기 때문에 (find)
• 분리집합 알고리즘으로 해결할 수 있다.
• 먼저 부모 노드들을 담을 parent(p) 리스트를 생성한다.
  • 초기 부모 노드는 각 집합에 담겨있는 숫자의 번호로 초기화한다.
  • list(range(n+1)) 코드를 통해 0부터 n까지 담겨있는 리스트를 쉽게 생성할 수 있다.
•  속해 있는 집합의 대표 번호(부모 노드 p)를 알기 위해서 find 함수를 정의한다.
  • 만약 부모 노드(p[x])가 원소 번호(x)와 같다면 자기 자신이 대표하는 노드라는 뜻으로 x를 return
  • 아니라면, 반복적으로 부모 노드를 찾아서 return 한다. 
• 같은 집합에 포함 시키기 위해 부모 노드를 일치시키는 union 함수를 정의한다.
  • a와 b의 부모 노드를 각각 찾고, 더 작은 노드를 부모 노드로 갖도록 한다.
  • 예를 들어 p[a] < p[b] 라면 p[p[b]] = p[a] 가 되는 것이다.
• 최종적으로 0번 연산은 union 함수를 실행시키고, 1번 연산은 find 함수를 통해 a와 b가 같은 집합에 포함되어 있다면 yes, 아니라면 no를 출력한다.

Python 코드

import sys
sys.setrecursionlimit(1000000) # 재귀 한계 늘려줌
input = sys.stdin.readline # 인풋 속도 개선

n,m = map(int,input().split())
p = list(range(n+1)) # 같은 집합인지 표시 (parent)

# 속한 집합 번호 출력
def find(x):
    if p[x] == x:
        return x
    p[x] = find(p[x])
    return p[x]

# 합집합, 즉 같은 집합 번호로 표기
def union(x,y):
    px = find(x)
    py = find(y)
    if x <= y:
        p[py] = px
    else:
        p[px] = py
        
def sol():
    res = [] # 정답을 모아서 한꺼번에 출력
    for _ in range(m):
        c,a,b = map(int,input().split())
        if c == 0: # 합집합 연산
            union(a,b)
        else:
            if find(a) == find(b):
                res.append("YES")
            else:
                res.append("NO")
    print(*res,sep="\n")

sol()

 

문제출처

1717번: 집합의 표현