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

문제

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

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

입력

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

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

예제 입력 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

예제 출력 1

NO
NO
YES

나의 풀이

이 문제는 유니온 파인드를 학습하기에 좋은 기본적인 문제로 유니온 파인드에 대한 이해를 필요로하는 문제였다.

• 유니온 파인드 (Union-Find)

  • 서로소 집합에 대한 정보를 저장하기에 좋은 자료구조로 union 함수와 find()함수를 이용하여 구현할 수 있다.
  • Tree구조를 이용하여 각 집합의 부모의 정보를 저장하고 있는다. ( 만약 Root라면 부모의 정보로 자신을 가지고 있게 된다. )
  • Find(a)
    • a집합의 부모의 정보를 사용하여 Root값을 찾아주는 함수이다.
  • Union(a,b)
    • a와 b집합이 서로소라면 합집합 연산을 하는 함수
    • 여기서 서로소인지 확인하기 위해서 find()함수를 사용하여 서로 root가 같은지를 확인한다.

  

  • 왼쪽그림은 각 집합이 Tree의 형식으로 연결이 되는 과정이고, 오른쪽 그림은 각 요소의 부모를 저장하는 배열을 단계별로 나타냈다.

위의 유니온 파인드를 이해했다면 그림에서 오른쪽에서 나타낸 것과 같이 각 요소에 부모에 대한 정보를 저장하여 이 문제를 해결할 수 있다. 처음에는 모두 서로소이므로 모든 원소가 자신의 인덱스를 가지게 된다. 그리고 앞에 0이 입력되어 합집합 연산을 진행해야하는 경우 union을 실행해주면 된다. 그리고 1이 입력되어 두 원소가 같은 집합에 속하는지를 확인하기 위해서는 find()함수를 사용하여 두 원소의 Root 원소를 찾아내고 같다면 같은 집합에 있는 것이고 그렇지않으면 서로소라는 것을 알 수 있다.

코드

# 1717번 집합의 표현
import sys

# find
def find(x):
    global parents

    while True:
        if parents[x] == x:
            return x
        else:
            x = parents[x]

# union
def union(x, y):
    global parents

    x = find(x)
    y = find(y)

    if not x == y:
        parents[y] = x

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

for _ in range(m):
    op, a, b = map(int, sys.stdin.readline().split())

    if op == 0:
        union(a,b)
    else:
        if (find(a) == find(b)):
            print("yes")
        else:
            print("no")